package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/keepmind9/acp-sdk-go/rpc"
	"github.com/keepmind9/acp-sdk-go/schema"
)

func ptrVers(v int) *schema.ProtocolVersion                                 { pv := schema.ProtocolVersion(v); return &pv }
func ptrStopReason(s schema.StopReason) *schema.StopReason                  { return &s }
func ptrOptId(s string) *schema.PermissionOptionId                          { o := schema.PermissionOptionId(s); return &o }
func ptrOptKind(k schema.PermissionOptionKind) *schema.PermissionOptionKind { return &k }
func ptrModeId(s string) *schema.SessionModeId                              { mid := schema.SessionModeId(s); return &mid }
func ptrConfigId(s string) *schema.SessionConfigId                          { cid := schema.SessionConfigId(s); return &cid }

// mockAgent implements Agent for testing.
type mockAgent struct {
	initResp    *schema.InitializeResponse
	sessionResp *schema.NewSessionResponse
	promptResp  *schema.PromptResponse
	cancelled   bool
}

func (m *mockAgent) Initialize(_ context.Context, req *schema.InitializeRequest) (*schema.InitializeResponse, error) {
	return m.initResp, nil
}

func (m *mockAgent) Authenticate(_ context.Context, req *schema.AuthenticateRequest) (*schema.AuthenticateResponse, error) {
	return &schema.AuthenticateResponse{}, nil
}

func (m *mockAgent) NewSession(_ context.Context, req *schema.NewSessionRequest) (*schema.NewSessionResponse, error) {
	return m.sessionResp, nil
}

func (m *mockAgent) LoadSession(_ context.Context, req *schema.LoadSessionRequest) (*schema.LoadSessionResponse, error) {
	return &schema.LoadSessionResponse{}, nil
}

func (m *mockAgent) ListSessions(_ context.Context, req *schema.ListSessionsRequest) (*schema.ListSessionsResponse, error) {
	return &schema.ListSessionsResponse{}, nil
}

func (m *mockAgent) Prompt(_ context.Context, req *schema.PromptRequest) (*schema.PromptResponse, error) {
	return m.promptResp, nil
}

func (m *mockAgent) Cancel(_ context.Context, notif *schema.CancelNotification) error {
	m.cancelled = true
	return nil
}

func (m *mockAgent) SetSessionMode(_ context.Context, req *schema.SetSessionModeRequest) (*schema.SetSessionModeResponse, error) {
	return &schema.SetSessionModeResponse{}, nil
}

func (m *mockAgent) SetSessionConfigOption(_ context.Context, req *schema.SetSessionConfigOptionRequest) (*schema.SetSessionConfigOptionResponse, error) {
	return &schema.SetSessionConfigOptionResponse{}, nil
}

// --- Router tests ---

func TestBuildAgentRouterInitialize(t *testing.T) {
	agent := &mockAgent{
		initResp: &schema.InitializeResponse{
			ProtocolVersion: ptrVers(1),
			AgentInfo:       &schema.Implementation{Name: "test-agent", Version: "0.1.0"},
		},
	}

	handler := buildAgentRouter(agent)
	params, _ := json.Marshal(schema.InitializeRequest{
		ProtocolVersion: ptrVers(1),
	})

	result, err := handler(schema.MethodInitialize, params, true)
	require.NoError(t, err)
	require.NotNil(t, result)

	resp, ok := result.(*schema.InitializeResponse)
	require.True(t, ok)
	require.NotNil(t, resp.ProtocolVersion)
	assert.Equal(t, 1, int(*resp.ProtocolVersion))
	require.NotNil(t, resp.AgentInfo)
	assert.Equal(t, "test-agent", resp.AgentInfo.Name)
}

func TestBuildAgentRouterNewSession(t *testing.T) {
	sid := schema.SessionId("sess-123")
	agent := &mockAgent{
		sessionResp: &schema.NewSessionResponse{SessionId: &sid},
	}
	handler := buildAgentRouter(agent)
	params, _ := json.Marshal(schema.NewSessionRequest{Cwd: "/tmp"})

	result, err := handler(schema.MethodSessionNew, params, true)
	require.NoError(t, err)
	resp, ok := result.(*schema.NewSessionResponse)
	require.True(t, ok)
	require.NotNil(t, resp.SessionId)
	assert.Equal(t, "sess-123", string(*resp.SessionId))
}

func TestBuildAgentRouterPrompt(t *testing.T) {
	sr := schema.StopReasonEndTurn
	agent := &mockAgent{
		promptResp: &schema.PromptResponse{StopReason: &sr},
	}
	handler := buildAgentRouter(agent)
	params, _ := json.Marshal(schema.PromptRequest{
		SessionId: ptrSessionId("sess-1"),
		Prompt:    []*schema.ContentBlock{{Type: "text", Text: &schema.TextContent{Text: "hello"}}},
	})

	result, err := handler(schema.MethodSessionPrompt, params, true)
	require.NoError(t, err)
	resp, ok := result.(*schema.PromptResponse)
	require.True(t, ok)
	require.NotNil(t, resp.StopReason)
	assert.Equal(t, schema.StopReasonEndTurn, *resp.StopReason)
}

func TestBuildAgentRouterSetSessionMode(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)
	params, _ := json.Marshal(schema.SetSessionModeRequest{
		SessionId: ptrSessionId("sess-1"),
		ModeId:    ptrModeId("plan"),
	})

	result, err := handler(schema.MethodSessionSet_mode, params, true)
	require.NoError(t, err)
	_, ok := result.(*schema.SetSessionModeResponse)
	require.True(t, ok)
}

func TestBuildAgentRouterSetConfigOption(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)
	params, _ := json.Marshal(schema.SetSessionConfigOptionRequest{
		SessionId: ptrSessionId("sess-1"),
		ConfigId:  ptrConfigId("opt-1"),
	})

	result, err := handler(schema.MethodSessionSet_config_option, params, true)
	require.NoError(t, err)
	_, ok := result.(*schema.SetSessionConfigOptionResponse)
	require.True(t, ok)
}

func TestBuildAgentRouterAuthenticate(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)
	params, _ := json.Marshal(schema.AuthenticateRequest{MethodId: "method-1"})

	result, err := handler(schema.MethodAuthenticate, params, true)
	require.NoError(t, err)
	_, ok := result.(*schema.AuthenticateResponse)
	require.True(t, ok)
}

func TestBuildAgentRouterLoadSession(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)
	params, _ := json.Marshal(schema.LoadSessionRequest{
		SessionId: ptrSessionId("sess-1"),
		Cwd:       "/tmp",
	})

	result, err := handler(schema.MethodSessionLoad, params, true)
	require.NoError(t, err)
	_, ok := result.(*schema.LoadSessionResponse)
	require.True(t, ok)
}

func TestBuildAgentRouterListSessions(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)
	params, _ := json.Marshal(schema.ListSessionsRequest{})

	result, err := handler(schema.MethodSessionList, params, true)
	require.NoError(t, err)
	_, ok := result.(*schema.ListSessionsResponse)
	require.True(t, ok)
}

func TestBuildAgentRouterCancel(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)

	params, _ := json.Marshal(schema.CancelNotification{SessionId: ptrSessionId("sess-1")})
	_, err := handler(schema.MethodSessionCancel, params, false)
	require.NoError(t, err)
	assert.True(t, agent.cancelled)
}

func TestBuildAgentRouterCancelAsRequest(t *testing.T) {
	// cancel is a notification, not a request
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)

	result, err := handler(schema.MethodSessionCancel, nil, true)
	require.NoError(t, err)
	assert.Nil(t, result)
}

func TestBuildAgentRouterUnknownMethod(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)

	_, err := handler("unknown/method", nil, true)
	require.Error(t, err)
	rpcErr, ok := err.(*rpc.RPCError)
	require.True(t, ok)
	assert.Equal(t, -32601, rpcErr.Code)
}

func TestBuildAgentRouterExtension(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)

	// Extension without handler -> method not found
	_, err := handler("_custom/method", nil, true)
	require.Error(t, err)
	rpcErr, ok := err.(*rpc.RPCError)
	require.True(t, ok)
	assert.Equal(t, -32601, rpcErr.Code)

	// Extension notification without handler -> silently ignored
	_, err = handler("_custom/notif", nil, false)
	assert.NoError(t, err)
}

func TestBuildAgentRouterInvalidParams(t *testing.T) {
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)

	_, err := handler(schema.MethodInitialize, json.RawMessage(`invalid`), true)
	require.Error(t, err)
	rpcErr, ok := err.(*rpc.RPCError)
	require.True(t, ok)
	assert.Equal(t, -32602, rpcErr.Code)
}

func TestBuildAgentRouterNotification(t *testing.T) {
	// Notifications (isRequest=false) for methods that are requests should be ignored
	agent := &mockAgent{}
	handler := buildAgentRouter(agent)

	params, _ := json.Marshal(schema.InitializeRequest{})
	result, err := handler(schema.MethodInitialize, params, false)
	require.NoError(t, err)
	assert.Nil(t, result)
}

// --- Connection tests ---

func TestAgentSideConnectionNotification(t *testing.T) {
	var buf bytes.Buffer
	agent := &mockAgent{}

	conn := NewAgentSideConnection(agent, nil, &buf)
	defer conn.Close()

	update := schema.SessionUpdate{
		SessionUpdate:     "agent_message_chunk",
		AgentMessageChunk: &schema.ContentChunk{Content: &schema.ContentBlock{Type: "text", Text: &schema.TextContent{Text: "hello world"}}},
	}
	err := conn.SessionUpdate("sess-1", &update)
	require.NoError(t, err)

	var msg map[string]any
	require.NoError(t, json.Unmarshal(bytes.TrimSpace(buf.Bytes()), &msg))
	assert.Equal(t, "2.0", msg["jsonrpc"])
	assert.Equal(t, "session/update", msg["method"])
}

func TestAgentSideConnectionRequestResponse(t *testing.T) {
	serverReadEnd, clientWriteEnd := io.Pipe()
	clientReadEnd, serverWriteEnd := io.Pipe()

	clientHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == schema.MethodSessionRequest_permission {
			return map[string]any{
				"outcome": "allow_once",
			}, nil
		}
		return nil, nil
	}

	clientConn := rpc.NewConnection(context.Background(), clientHandler, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	agent := &mockAgent{}
	agentConn := NewAgentSideConnection(agent, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	resp, err := agentConn.RequestPermission(&schema.RequestPermissionRequest{
		SessionId: ptrSessionId("sess-1"),
		Options: []*schema.PermissionOption{
			{OptionId: ptrOptId("opt-1"), Name: "Allow", Kind: ptrOptKind(schema.PermissionOptionKindAllowOnce)},
		},
	})
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.NotNil(t, resp.Outcome)
	assert.Equal(t, "allow_once", string(*resp.Outcome))

	serverReadEnd.Close()
	serverWriteEnd.Close()
	clientReadEnd.Close()
	clientWriteEnd.Close()
}

func TestAgentSideConnectionReadTextFile(t *testing.T) {
	serverReadEnd, clientWriteEnd := io.Pipe()
	clientReadEnd, serverWriteEnd := io.Pipe()

	clientHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == schema.MethodFsRead_text_file {
			return map[string]any{
				"content": "file contents",
			}, nil
		}
		return nil, nil
	}

	clientConn := rpc.NewConnection(context.Background(), clientHandler, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	agent := &mockAgent{}
	agentConn := NewAgentSideConnection(agent, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	resp, err := agentConn.ReadTextFile(&schema.ReadTextFileRequest{
		SessionId: ptrSessionId("sess-1"),
		Path:      "main.go",
	})
	require.NoError(t, err)
	require.NotNil(t, resp)

	serverReadEnd.Close()
	serverWriteEnd.Close()
	clientReadEnd.Close()
	clientWriteEnd.Close()
}

func TestAgentSideConnectionWriteTextFile(t *testing.T) {
	serverReadEnd, clientWriteEnd := io.Pipe()
	clientReadEnd, serverWriteEnd := io.Pipe()

	clientHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == schema.MethodFsWrite_text_file {
			return map[string]any{"success": true}, nil
		}
		return nil, nil
	}

	clientConn := rpc.NewConnection(context.Background(), clientHandler, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	agent := &mockAgent{}
	agentConn := NewAgentSideConnection(agent, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	_, err := agentConn.WriteTextFile(&schema.WriteTextFileRequest{
		SessionId: ptrSessionId("sess-1"),
		Path:      "main.go",
		Content:   "package main",
	})
	require.NoError(t, err)

	serverReadEnd.Close()
	serverWriteEnd.Close()
	clientReadEnd.Close()
	clientWriteEnd.Close()
}

func TestAgentSideConnectionTerminal(t *testing.T) {
	serverReadEnd, clientWriteEnd := io.Pipe()
	clientReadEnd, serverWriteEnd := io.Pipe()

	clientHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == schema.MethodTerminalCreate {
			return map[string]any{"terminalId": "term-1"}, nil
		}
		return nil, nil
	}

	clientConn := rpc.NewConnection(context.Background(), clientHandler, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	agent := &mockAgent{}
	agentConn := NewAgentSideConnection(agent, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	resp, err := agentConn.CreateTerminal(&schema.CreateTerminalRequest{
		SessionId: ptrSessionId("sess-1"),
		Command:   "bash",
	})
	require.NoError(t, err)
	require.NotNil(t, resp)

	serverReadEnd.Close()
	serverWriteEnd.Close()
	clientReadEnd.Close()
	clientWriteEnd.Close()
}

func TestAgentSideConnectionExtMethod(t *testing.T) {
	serverReadEnd, clientWriteEnd := io.Pipe()
	clientReadEnd, serverWriteEnd := io.Pipe()

	clientHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == "_custom" {
			return map[string]any{"result": "custom response"}, nil
		}
		return nil, nil
	}

	clientConn := rpc.NewConnection(context.Background(), clientHandler, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	agent := &mockAgent{}
	agentConn := NewAgentSideConnection(agent, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	result, err := agentConn.ExtMethod("custom", map[string]any{"key": "value"})
	require.NoError(t, err)
	require.NotNil(t, result)

	serverReadEnd.Close()
	serverWriteEnd.Close()
	clientReadEnd.Close()
	clientWriteEnd.Close()
}

func TestAgentSideConnectionExtNotification(t *testing.T) {
	var buf bytes.Buffer
	agent := &mockAgent{}

	conn := NewAgentSideConnection(agent, nil, &buf)
	defer conn.Close()

	err := conn.ExtNotification("custom", map[string]any{"key": "value"})
	require.NoError(t, err)

	var msg map[string]any
	require.NoError(t, json.Unmarshal(bytes.TrimSpace(buf.Bytes()), &msg))
	assert.Equal(t, "2.0", msg["jsonrpc"])
	assert.Equal(t, "_custom", msg["method"])
}

func TestAgentSideConnectionClose(t *testing.T) {
	r, w := io.Pipe()
	defer r.Close()
	defer w.Close()

	agent := &mockAgent{}
	conn := NewAgentSideConnection(agent, r, w)
	require.NoError(t, conn.Close())
}