package acp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"strings"
	"sync"
	"testing"
	"time"
)

// stubHandler answers the core methods with whatever the test set on it.
type stubHandler struct {
	initialize func(InitializeRequest) (InitializeResponse, error)
	authFail   error
	newSession func(NewSessionRequest) (NewSessionResponse, error)
	prompt     func(context.Context, PromptRequest) (PromptResponse, error)

	mu        sync.Mutex
	cancelled []string
}

func (s *stubHandler) Initialize(_ context.Context, req InitializeRequest) (InitializeResponse, error) {
	if s.initialize != nil {
		return s.initialize(req)
	}
	return InitializeResponse{ProtocolVersion: ProtocolVersion, AuthMethods: []AuthMethod{}}, nil
}

func (s *stubHandler) Authenticate(context.Context, AuthenticateRequest) error { return s.authFail }

func (s *stubHandler) NewSession(_ context.Context, req NewSessionRequest) (NewSessionResponse, error) {
	if s.newSession != nil {
		return s.newSession(req)
	}
	return NewSessionResponse{SessionID: "s1"}, nil
}

func (s *stubHandler) Prompt(ctx context.Context, req PromptRequest) (PromptResponse, error) {
	if s.prompt != nil {
		return s.prompt(ctx, req)
	}
	return PromptResponse{StopReason: StopEndTurn}, nil
}

func (s *stubHandler) Cancel(_ context.Context, p CancelParams) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.cancelled = append(s.cancelled, p.SessionID)
}

// peer is the far end of a connection: a pipe to write requests into and a
// reader that hands back one decoded frame at a time.
type peer struct {
	t    *testing.T
	in   *io.PipeWriter
	out  *bufio.Reader
	conn *Conn
	done chan error
}

func newPeer(t *testing.T, h Handler) *peer {
	t.Helper()
	toAgent, fromPeer := io.Pipe()
	fromAgent, toPeer := io.Pipe()

	c := NewConn(toAgent, toPeer, h)
	p := &peer{t: t, in: fromPeer, out: bufio.NewReader(fromAgent), conn: c, done: make(chan error, 1)}

	go func() { p.done <- c.Serve(context.Background()) }()
	t.Cleanup(func() {
		_ = fromPeer.Close()
		select {
		case <-p.done:
		case <-time.After(2 * time.Second):
			t.Error("Serve did not return after the stream closed")
		}
	})
	return p
}

func (p *peer) send(line string) {
	p.t.Helper()
	if _, err := io.WriteString(p.in, line+"\n"); err != nil {
		p.t.Fatalf("writing to the agent: %v", err)
	}
}

// readRaw returns the next line exactly as it went over the wire, because some
// of what matters here — an absent key — does not survive decoding.
func (p *peer) readRaw() string {
	p.t.Helper()
	type result struct {
		line string
		err  error
	}
	ch := make(chan result, 1)
	go func() {
		line, err := p.out.ReadString('\n')
		ch <- result{line, err}
	}()
	select {
	case r := <-ch:
		if r.err != nil {
			p.t.Fatalf("reading from the agent: %v", r.err)
		}
		return strings.TrimSuffix(r.line, "\n")
	case <-time.After(2 * time.Second):
		p.t.Fatal("timed out waiting for a frame from the agent")
		return ""
	}
}

func (p *peer) read() Message {
	p.t.Helper()
	var m Message
	line := p.readRaw()
	if err := json.Unmarshal([]byte(line), &m); err != nil {
		p.t.Fatalf("undecodable frame %q: %v", line, err)
	}
	return m
}

func TestEveryFrameCarriesTheProtocolVersion(t *testing.T) {
	p := newPeer(t, &stubHandler{})

	p.send(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":1}}`)
	if got := p.read().JSONRPC; got != "2.0" {
		t.Errorf("response jsonrpc = %q, want 2.0", got)
	}

	go func() {
		_ = p.conn.Notify(context.Background(), MethodSessionUpdate, SessionNotification{SessionID: "s1"})
	}()
	if got := p.read().JSONRPC; got != "2.0" {
		t.Errorf("notification jsonrpc = %q, want 2.0", got)
	}
}

func TestANotificationHasNoIDKeyAtAll(t *testing.T) {
	p := newPeer(t, &stubHandler{})

	go func() {
		_ = p.conn.Notify(context.Background(), MethodSessionUpdate, SessionNotification{
			SessionID: "s1",
			Update:    Update{SessionUpdate: UpdateAgentMessageChunk, Content: TextContent("hi")},
		})
	}()

	// A decoded map would show the difference between absent and null as the
	// same missing key, and null is the thing that breaks the client.
	line := p.readRaw()
	if strings.Contains(line, `"id"`) {
		t.Errorf("notification carries an id: %s", line)
	}
	if !strings.Contains(line, `"sessionUpdate":"agent_message_chunk"`) {
		t.Errorf("notification lost its tag: %s", line)
	}
}

func TestUnknownMethodIsMethodNotFound(t *testing.T) {
	p := newPeer(t, &stubHandler{})

	p.send(`{"jsonrpc":"2.0","id":"abc","method":"session/set_model","params":{}}`)
	m := p.read()

	if m.ID == nil || string(*m.ID) != `"abc"` {
		t.Errorf("id did not come back byte for byte: %v", m.ID)
	}
	if m.Error == nil || m.Error.Code != CodeMethodNotFound {
		t.Fatalf("want a method-not-found error, got %+v", m.Error)
	}
}

// extendedHandler answers anything the core set does not.
type extendedHandler struct {
	stubHandler
	seen string
}

func (e *extendedHandler) Extra(_ context.Context, method string, _ json.RawMessage) (any, error) {
	e.seen = method
	return map[string]any{}, nil
}

func TestAnExtenderAnswersTheRest(t *testing.T) {
	h := &extendedHandler{}
	p := newPeer(t, h)

	p.send(`{"jsonrpc":"2.0","id":7,"method":"session/set_model","params":{"modelId":"x"}}`)
	m := p.read()

	if m.Error != nil {
		t.Fatalf("unexpected error: %v", m.Error)
	}
	if h.seen != "session/set_model" {
		t.Errorf("Extra saw %q", h.seen)
	}
	if string(m.Result) != "{}" {
		t.Errorf("result = %s, want {}", m.Result)
	}
}

func TestAHandlerErrorBecomesInternalAndTheLoopSurvives(t *testing.T) {
	h := &stubHandler{initialize: func(InitializeRequest) (InitializeResponse, error) {
		return InitializeResponse{}, errors.New("no")
	}}
	p := newPeer(t, h)

	p.send(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`)
	if m := p.read(); m.Error == nil || m.Error.Code != CodeInternal {
		t.Fatalf("want an internal error, got %+v", m.Error)
	}

	// The connection is still usable, which is the point.
	p.send(`{"jsonrpc":"2.0","id":2,"method":"session/new","params":{"cwd":"/tmp"}}`)
	if m := p.read(); m.Error != nil {
		t.Fatalf("the connection did not survive the error: %v", m.Error)
	}
}

func TestBadParamsAreInvalidParamsNotInternal(t *testing.T) {
	p := newPeer(t, &stubHandler{})

	p.send(`{"jsonrpc":"2.0","id":1,"method":"session/new","params":{"cwd":42}}`)
	if m := p.read(); m.Error == nil || m.Error.Code != CodeInvalidParams {
		t.Fatalf("want invalid params, got %+v", m.Error)
	}
}

func TestMalformedLinesAreAnsweredWhenTheyCanBeAndSkippedWhenTheyCannot(t *testing.T) {
	p := newPeer(t, &stubHandler{})

	// Unparseable as a whole, but the id is recoverable: the peer is waiting.
	p.send(`{"jsonrpc":"2.0","id":3,"method":"initialize","params":{`)
	if m := p.read(); m.Error == nil || m.Error.Code != CodeParse {
		t.Fatalf("want a parse error, got %+v", m.Error)
	}

	// Nothing to answer, and nothing to do but carry on.
	p.send(`not json at all`)
	p.send(`{"jsonrpc":"2.0","id":4,"method":"initialize","params":{}}`)
	if m := p.read(); m.Error != nil {
		t.Fatalf("the loop did not survive a junk line: %v", m.Error)
	}
}

func TestRequestWaitsForTheMatchingResponse(t *testing.T) {
	p := newPeer(t, &stubHandler{})

	var out RequestPermissionResult
	errc := make(chan error, 1)
	go func() {
		errc <- p.conn.Request(context.Background(), MethodRequestPermission,
			RequestPermissionParams{SessionID: "s1"}, &out)
	}()

	req := p.read()
	if req.Method != MethodRequestPermission {
		t.Fatalf("method = %q", req.Method)
	}
	if req.ID == nil {
		t.Fatal("an outbound request must carry an id")
	}
	p.send(`{"jsonrpc":"2.0","id":` + string(*req.ID) + `,"result":{"outcome":{"outcome":"selected","optionId":"allow-once"}}}`)

	if err := <-errc; err != nil {
		t.Fatalf("Request: %v", err)
	}
	if out.Outcome.OptionID != OptAllowOnce {
		t.Errorf("outcome = %+v", out.Outcome)
	}
}

func TestRequestGivesUpWithItsContext(t *testing.T) {
	p := newPeer(t, &stubHandler{})

	ctx, cancel := context.WithCancel(context.Background())
	errc := make(chan error, 1)
	go func() {
		errc <- p.conn.Request(ctx, MethodRequestPermission, RequestPermissionParams{}, nil)
	}()
	p.read() // let it go out before cancelling
	cancel()

	select {
	case err := <-errc:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("want context.Canceled, got %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Request did not return when its context was cancelled")
	}
}

func TestRequestFailsWhenTheConnectionEnds(t *testing.T) {
	toAgent, fromPeer := io.Pipe()
	fromAgent, toPeer := io.Pipe()
	// Drained so the outbound request cannot block on an unread pipe.
	go func() { _, _ = io.Copy(io.Discard, fromAgent) }()

	c := NewConn(toAgent, toPeer, &stubHandler{})
	served := make(chan error, 1)
	go func() { served <- c.Serve(context.Background()) }()

	errc := make(chan error, 1)
	go func() {
		errc <- c.Request(context.Background(), MethodRequestPermission, RequestPermissionParams{}, nil)
	}()

	time.Sleep(20 * time.Millisecond)
	_ = fromPeer.Close()

	select {
	case err := <-errc:
		if !errors.Is(err, ErrClosed) {
			t.Fatalf("want ErrClosed, got %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("a pending request outlived its connection")
	}
	if err := <-served; err != nil {
		t.Errorf("Serve after a clean close: %v", err)
	}
}

func TestCancelArrivesWhileAPromptIsStillRunning(t *testing.T) {
	release := make(chan struct{})
	h := &stubHandler{}
	h.prompt = func(ctx context.Context, _ PromptRequest) (PromptResponse, error) {
		<-release
		return PromptResponse{StopReason: StopCancelled}, nil
	}
	p := newPeer(t, h)

	p.send(`{"jsonrpc":"2.0","id":1,"method":"session/prompt","params":{"sessionId":"s1","prompt":[]}}`)
	p.send(`{"jsonrpc":"2.0","method":"session/cancel","params":{"sessionId":"s1"}}`)

	// The cancel must land while the prompt is still blocked. Polling rather
	// than sleeping a fixed time keeps the test honest on a loaded machine.
	deadline := time.Now().Add(2 * time.Second)
	for {
		h.mu.Lock()
		got := len(h.cancelled)
		h.mu.Unlock()
		if got == 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("the cancel was queued behind the prompt")
		}
		time.Sleep(time.Millisecond)
	}

	close(release)
	if m := p.read(); m.Error != nil {
		t.Fatalf("prompt failed: %v", m.Error)
	}
}

func TestCleanEOFIsNotAnError(t *testing.T) {
	c := NewConn(strings.NewReader(""), &bytes.Buffer{}, &stubHandler{})
	if err := c.Serve(context.Background()); err != nil {
		t.Errorf("Serve on an empty stream: %v", err)
	}
}

func TestATrailingLineWithoutANewlineIsStillHandled(t *testing.T) {
	// A peer that exits right after writing may not terminate its last line.
	var out bytes.Buffer
	c := NewConn(strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`), &out, &stubHandler{})
	if err := c.Serve(context.Background()); err != nil {
		t.Fatalf("Serve: %v", err)
	}
	if !strings.Contains(out.String(), `"protocolVersion":1`) {
		t.Errorf("the last line was dropped: %q", out.String())
	}
}