package client

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"
)

// mockClient implements Client for testing.
type mockClient struct {
	lastSessionNotif *schema.SessionNotification
	permResp         *schema.RequestPermissionResponse
	readResp         *schema.ReadTextFileResponse
	writeResp        *schema.WriteTextFileResponse
	terminalResp     *schema.CreateTerminalResponse
}

func (m *mockClient) RequestPermission(_ context.Context, req *schema.RequestPermissionRequest) (*schema.RequestPermissionResponse, error) {
	return m.permResp, nil
}

func (m *mockClient) SessionUpdate(_ context.Context, notif *schema.SessionNotification) error {
	m.lastSessionNotif = notif
	return nil
}

func (m *mockClient) ReadTextFile(_ context.Context, req *schema.ReadTextFileRequest) (*schema.ReadTextFileResponse, error) {
	return m.readResp, nil
}

func (m *mockClient) WriteTextFile(_ context.Context, req *schema.WriteTextFileRequest) (*schema.WriteTextFileResponse, error) {
	return m.writeResp, nil
}

func (m *mockClient) CreateTerminal(_ context.Context, req *schema.CreateTerminalRequest) (*schema.CreateTerminalResponse, error) {
	return m.terminalResp, nil
}

func (m *mockClient) TerminalOutput(_ context.Context, req *schema.TerminalOutputRequest) (*schema.TerminalOutputResponse, error) {
	return &schema.TerminalOutputResponse{}, nil
}

func (m *mockClient) ReleaseTerminal(_ context.Context, req *schema.ReleaseTerminalRequest) (*schema.ReleaseTerminalResponse, error) {
	return &schema.ReleaseTerminalResponse{}, nil
}

func (m *mockClient) WaitForTerminalExit(_ context.Context, req *schema.WaitForTerminalExitRequest) (*schema.WaitForTerminalExitResponse, error) {
	return &schema.WaitForTerminalExitResponse{}, nil
}

func (m *mockClient) KillTerminal(_ context.Context, req *schema.KillTerminalRequest) (*schema.KillTerminalResponse, error) {
	return &schema.KillTerminalResponse{}, nil
}

// --- Router tests ---

func TestBuildClientRouterSessionUpdate(t *testing.T) {
	mc := &mockClient{}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.SessionNotification{
		SessionId: ptrSessionId("sess-1"),
		Update: &schema.SessionUpdate{
			SessionUpdate:     "agent_message_chunk",
			AgentMessageChunk: &schema.ContentChunk{Content: &schema.ContentBlock{Type: "text", Text: &schema.TextContent{Text: "hello"}}},
		},
	})

	_, err := handler(schema.MethodSessionUpdate, params, false)
	require.NoError(t, err)
	require.NotNil(t, mc.lastSessionNotif)
	require.NotNil(t, mc.lastSessionNotif.SessionId)
	assert.Equal(t, "sess-1", string(*mc.lastSessionNotif.SessionId))
}

func TestBuildClientRouterRequestPermission(t *testing.T) {
	outcome := schema.RequestPermissionOutcomeSelected
	mc := &mockClient{
		permResp: &schema.RequestPermissionResponse{
			Outcome: &outcome,
		},
	}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.RequestPermissionRequest{
		SessionId: ptrSessionId("sess-1"),
	})
	result, err := handler(schema.MethodSessionRequest_permission, params, true)
	require.NoError(t, err)
	require.NotNil(t, result)
}

func TestBuildClientRouterReadTextFile(t *testing.T) {
	mc := &mockClient{
		readResp: &schema.ReadTextFileResponse{Content: "file content"},
	}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.ReadTextFileRequest{
		SessionId: ptrSessionId("sess-1"),
		Path:      "main.go",
	})
	result, err := handler(schema.MethodFsRead_text_file, params, true)
	require.NoError(t, err)
	require.NotNil(t, result)
}

func TestBuildClientRouterWriteTextFile(t *testing.T) {
	mc := &mockClient{
		writeResp: &schema.WriteTextFileResponse{},
	}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.WriteTextFileRequest{
		SessionId: ptrSessionId("sess-1"),
		Path:      "main.go",
		Content:   "package main",
	})
	_, err := handler(schema.MethodFsWrite_text_file, params, true)
	require.NoError(t, err)
}

func TestBuildClientRouterCreateTerminal(t *testing.T) {
	mc := &mockClient{
		terminalResp: &schema.CreateTerminalResponse{TerminalId: "term-1"},
	}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.CreateTerminalRequest{
		SessionId: ptrSessionId("sess-1"),
		Command:   "bash",
	})
	result, err := handler(schema.MethodTerminalCreate, params, true)
	require.NoError(t, err)
	require.NotNil(t, result)
}

func TestBuildClientRouterTerminalOutput(t *testing.T) {
	mc := &mockClient{}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.TerminalOutputRequest{
		SessionId:  ptrSessionId("sess-1"),
		TerminalId: "term-1",
	})
	result, err := handler(schema.MethodTerminalOutput, params, true)
	require.NoError(t, err)
	require.NotNil(t, result)
}

func TestBuildClientRouterReleaseTerminal(t *testing.T) {
	mc := &mockClient{}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.ReleaseTerminalRequest{
		SessionId:  ptrSessionId("sess-1"),
		TerminalId: "term-1",
	})
	_, err := handler(schema.MethodTerminalRelease, params, true)
	require.NoError(t, err)
}

func TestBuildClientRouterWaitForTerminalExit(t *testing.T) {
	mc := &mockClient{}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.WaitForTerminalExitRequest{
		SessionId:  ptrSessionId("sess-1"),
		TerminalId: "term-1",
	})
	result, err := handler(schema.MethodTerminalWait_for_exit, params, true)
	require.NoError(t, err)
	require.NotNil(t, result)
}

func TestBuildClientRouterKillTerminal(t *testing.T) {
	mc := &mockClient{}
	handler := buildClientRouter(mc)

	params, _ := json.Marshal(schema.KillTerminalRequest{
		SessionId:  ptrSessionId("sess-1"),
		TerminalId: "term-1",
	})
	_, err := handler(schema.MethodTerminalKill, params, true)
	require.NoError(t, err)
}

func TestBuildClientRouterUnknownMethod(t *testing.T) {
	mc := &mockClient{}
	handler := buildClientRouter(mc)

	_, 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 TestBuildClientRouterInvalidParams(t *testing.T) {
	mc := &mockClient{}
	handler := buildClientRouter(mc)

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

func TestBuildClientRouterExtension(t *testing.T) {
	mc := &mockClient{}
	handler := buildClientRouter(mc)

	_, 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)

	_, err = handler("_custom/notif", nil, false)
	assert.NoError(t, err)
}

func TestBuildClientRouterNotification(t *testing.T) {
	// Methods that are requests should ignore notifications
	mc := &mockClient{}
	handler := buildClientRouter(mc)

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

// --- Connection tests ---

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

	agentHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == schema.MethodInitialize {
			return map[string]any{
				"protocolVersion": 1,
				"agentInfo":       map[string]string{"name": "test-agent", "version": "0.1.0"},
			}, nil
		}
		return nil, rpc.NewRPCError(-32601, "Method not found")
	}

	agentConn := rpc.NewConnection(context.Background(), agentHandler, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	mc := &mockClient{}
	clientConn := NewClientSideConnection(mc, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	resp, err := clientConn.Initialize(&schema.InitializeRequest{
		ProtocolVersion:    ptrVers(1),
		ClientCapabilities: &schema.ClientCapabilities{},
		ClientInfo:         &schema.Implementation{Name: "test-client", Version: "1.0.0"},
	})
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.NotNil(t, resp.ProtocolVersion)
	assert.Equal(t, 1, int(*resp.ProtocolVersion))

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

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

	agentHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == schema.MethodSessionNew {
			return map[string]any{"sessionId": "sess-123"}, nil
		}
		return nil, nil
	}

	agentConn := rpc.NewConnection(context.Background(), agentHandler, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	mc := &mockClient{}
	clientConn := NewClientSideConnection(mc, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	resp, err := clientConn.NewSession(&schema.NewSessionRequest{Cwd: "/tmp"})
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.NotNil(t, resp.SessionId)
	assert.Equal(t, "sess-123", string(*resp.SessionId))

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

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

	agentHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == schema.MethodSessionPrompt {
			return map[string]any{"stopReason": "end_turn"}, nil
		}
		return nil, nil
	}

	agentConn := rpc.NewConnection(context.Background(), agentHandler, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	mc := &mockClient{}
	clientConn := NewClientSideConnection(mc, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	resp, err := clientConn.Prompt(&schema.PromptRequest{
		SessionId: ptrSessionId("sess-1"),
		Prompt:    []*schema.ContentBlock{{Type: "text", Text: &schema.TextContent{Text: "hello"}}},
	})
	require.NoError(t, err)
	require.NotNil(t, resp)

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

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

	agentHandler := func(method string, params json.RawMessage, isRequest bool) (any, error) {
		if method == schema.MethodSessionSet_mode {
			return map[string]any{}, nil
		}
		return nil, nil
	}

	agentConn := rpc.NewConnection(context.Background(), agentHandler, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	mc := &mockClient{}
	clientConn := NewClientSideConnection(mc, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

	_, err := clientConn.SetSessionMode(&schema.SetSessionModeRequest{
		SessionId: ptrSessionId("sess-1"),
		ModeId:    ptrModeId("plan"),
	})
	require.NoError(t, err)

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

func TestClientSideConnectionCancel(t *testing.T) {
	var buf bytes.Buffer
	mc := &mockClient{}

	conn := NewClientSideConnection(mc, nil, &buf)
	defer conn.Close()

	err := conn.Cancel("sess-1")
	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/cancel", msg["method"])
}

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

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

	agentConn := rpc.NewConnection(context.Background(), agentHandler, serverReadEnd, serverWriteEnd)
	defer agentConn.Close()
	go agentConn.ReceiveLoop()

	mc := &mockClient{}
	clientConn := NewClientSideConnection(mc, clientReadEnd, clientWriteEnd)
	defer clientConn.Close()
	go clientConn.ReceiveLoop()

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

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

func TestClientSideConnectionExtNotification(t *testing.T) {
	var buf bytes.Buffer
	mc := &mockClient{}

	conn := NewClientSideConnection(mc, 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 TestClientSideConnectionClose(t *testing.T) {
	r, w := io.Pipe()
	defer r.Close()
	defer w.Close()

	mc := &mockClient{}
	conn := NewClientSideConnection(mc, r, w)
	require.NoError(t, conn.Close())
}