package serve

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"codey/internal/acp"
	"codey/internal/agent"
	"codey/internal/llm"
	"codey/internal/metrics"
	"codey/internal/tools"
)

// --- a fake Anthropic endpoint -------------------------------------------

type fakeAPI struct {
	t   *testing.T
	srv *httptest.Server

	mu        sync.Mutex
	responses []string
	requests  [][]byte
	n         int
	// hold, when non-nil, is waited on before the first response goes out, so
	// a test can cancel a turn that is genuinely in flight.
	hold chan struct{}
}

func newFakeAPI(t *testing.T, responses ...string) *fakeAPI {
	f := &fakeAPI{t: t, responses: responses}
	f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		f.mu.Lock()
		hold := f.hold
		f.mu.Unlock()
		if hold != nil {
			select {
			case <-hold:
			case <-r.Context().Done():
				return
			}
		}

		body, _ := io.ReadAll(r.Body)

		f.mu.Lock()
		i := f.n
		f.n++
		f.requests = append(f.requests, body)
		f.mu.Unlock()

		if i >= len(f.responses) {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		io.WriteString(w, f.responses[i])
	}))
	t.Cleanup(f.srv.Close)
	return f
}

func msg(stopReason string, blocks ...string) string {
	return `{"id":"msg","type":"message","role":"assistant","model":"test-model","content":[` +
		strings.Join(blocks, ",") + `],"stop_reason":"` + stopReason + `",` +
		`"usage":{"input_tokens":10,"output_tokens":5}}`
}

func textBlock(s string) string { return `{"type":"text","text":"` + s + `"}` }

func toolUseBlock(id, name, input string) string {
	return `{"type":"tool_use","id":"` + id + `","name":"` + name + `","input":` + input + `}`
}

// --- the server under test, wired to that endpoint ------------------------

func newTestServer(t *testing.T, f *fakeAPI, opts Options) *testClient {
	t.Helper()

	var logBuf bytes.Buffer
	log := metrics.NewWriter(&logBuf, "test")

	if opts.NewAgent == nil {
		opts.NewAgent = func(root, model string, state *agent.State) (*agent.Agent, error) {
			if model == "" {
				model = "test-model"
			}
			reg := tools.NewRegistry(log).Register(
				tools.NewRead(root, state),
				tools.NewWrite(root, state),
				tools.NewBash(root),
			)
			client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"), llm.WithModel(model))
			// Streaming needs an SSE endpoint; the delta path is covered in
			// update_test.go, and a whole-response turn exercises everything
			// this file is about.
			return agent.New(client, reg, state, log, agent.Options{}), nil
		}
	}
	if opts.Version == "" {
		opts.Version = "v0.0.0-test"
	}

	toAgent, fromClient := io.Pipe()
	fromAgent, toClient := io.Pipe()

	srv := New(opts)
	served := make(chan error, 1)
	go func() { served <- srv.Run(context.Background(), toAgent, toClient) }()

	c := &testClient{t: t, in: fromClient, frames: make(chan acp.Message, 128)}
	go c.read(bufio.NewReader(fromAgent))

	t.Cleanup(func() {
		_ = fromClient.Close()
		select {
		case <-served:
		case <-time.After(3 * time.Second):
			t.Error("Run did not return after the stream closed")
		}
	})
	return c
}

// testClient is the far end: it sends requests, answers the ones the agent
// sends back, and remembers every notification it saw.
type testClient struct {
	t      *testing.T
	in     io.WriteCloser
	frames chan acp.Message

	nextID int

	// onPermission answers a permission request. Nil means "reject", which is
	// what an unprepared test should get rather than a hang.
	onPermission func(acp.RequestPermissionParams) acp.RequestPermissionResult

	mu      sync.Mutex
	updates []acp.SessionNotification
	// raw is every line the agent wrote, kept so a test can assert on what
	// reached the stream rather than on what decoded off it.
	sent []string
	// permission is the last permission request, unparsed.
	permission json.RawMessage
}

// lines is every frame the agent wrote, in order.
func (c *testClient) lines() []string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return append([]string(nil), c.sent...)
}

func (c *testClient) read(br *bufio.Reader) {
	for {
		line, err := br.ReadBytes('\n')
		if trimmed := bytes.TrimSpace(line); len(trimmed) > 0 {
			c.mu.Lock()
			c.sent = append(c.sent, string(trimmed))
			c.mu.Unlock()

			var m acp.Message
			if json.Unmarshal(line, &m) == nil {
				c.frames <- m
			}
		}
		if err != nil {
			close(c.frames)
			return
		}
	}
}

func (c *testClient) write(m any) {
	c.t.Helper()
	b, err := json.Marshal(m)
	if err != nil {
		c.t.Fatal(err)
	}
	if _, err := c.in.Write(append(b, '\n')); err != nil {
		c.t.Fatalf("writing to the agent: %v", err)
	}
}

// call sends a request and pumps the connection until its answer arrives,
// recording notifications and answering anything the agent asks on the way.
func (c *testClient) call(method string, params any) acp.Message {
	c.t.Helper()
	c.nextID++
	id := json.RawMessage(json.Number(itoa(c.nextID)).String())
	c.write(acp.Message{JSONRPC: "2.0", ID: &id, Method: method, Params: mustJSON(c.t, params)})
	return c.pump(string(id))
}

func (c *testClient) notify(method string, params any) {
	c.t.Helper()
	c.write(acp.Message{JSONRPC: "2.0", Method: method, Params: mustJSON(c.t, params)})
}

func (c *testClient) pump(wantID string) acp.Message {
	c.t.Helper()
	deadline := time.After(10 * time.Second)
	for {
		select {
		case m, ok := <-c.frames:
			if !ok {
				c.t.Fatal("the connection closed while waiting for a response")
			}
			switch {
			case m.Method == acp.MethodSessionUpdate:
				var n acp.SessionNotification
				if err := json.Unmarshal(m.Params, &n); err == nil {
					c.mu.Lock()
					c.updates = append(c.updates, n)
					c.mu.Unlock()
				}
			case m.Method == acp.MethodRequestPermission:
				c.answerPermission(m)
			case m.ID != nil && string(*m.ID) == wantID:
				return m
			}
		case <-deadline:
			c.t.Fatalf("timed out waiting for the answer to request %s", wantID)
		}
	}
}

func (c *testClient) answerPermission(m acp.Message) {
	c.t.Helper()
	var params acp.RequestPermissionParams
	_ = json.Unmarshal(m.Params, &params)

	c.mu.Lock()
	c.permission = m.Params
	c.mu.Unlock()

	result := acp.RequestPermissionResult{
		Outcome: acp.PermissionOutcome{Outcome: acp.OutcomeSelected, OptionID: acp.OptRejectOnce},
	}
	if c.onPermission != nil {
		result = c.onPermission(params)
	}
	c.write(acp.Message{JSONRPC: "2.0", ID: m.ID, Result: mustJSON(c.t, result)})
}

// text is everything the agent said, across every session.
func (c *testClient) text() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	var b strings.Builder
	for _, n := range c.updates {
		if n.Update.SessionUpdate != acp.UpdateAgentMessageChunk {
			continue
		}
		var content struct {
			Text string `json:"text"`
		}
		raw, _ := json.Marshal(n.Update.Content)
		_ = json.Unmarshal(raw, &content)
		b.WriteString(content.Text)
	}
	return b.String()
}

func (c *testClient) sessionUpdates(id string) []acp.Update {
	c.mu.Lock()
	defer c.mu.Unlock()
	var out []acp.Update
	for _, n := range c.updates {
		if n.SessionID == id {
			out = append(out, n.Update)
		}
	}
	return out
}

func mustJSON(t *testing.T, v any) json.RawMessage {
	t.Helper()
	b, err := json.Marshal(v)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

// handshake runs initialize, authenticate and session/new, and returns the id.
func (c *testClient) handshake(cwd string) string {
	c.t.Helper()

	init := c.call(acp.MethodInitialize, acp.InitializeRequest{ProtocolVersion: 1})
	if init.Error != nil {
		c.t.Fatalf("initialize: %v", init.Error)
	}
	var initResp acp.InitializeResponse
	if err := json.Unmarshal(init.Result, &initResp); err != nil {
		c.t.Fatalf("initialize result: %v", err)
	}
	if initResp.ProtocolVersion != acp.ProtocolVersion {
		c.t.Fatalf("protocol version = %d", initResp.ProtocolVersion)
	}

	if m := c.call(acp.MethodAuthenticate, acp.AuthenticateRequest{MethodID: "cached_token"}); m.Error != nil {
		c.t.Fatalf("authenticate: %v", m.Error)
	}

	m := c.call(acp.MethodNewSession, acp.NewSessionRequest{Cwd: cwd})
	if m.Error != nil {
		c.t.Fatalf("session/new: %v", m.Error)
	}
	var resp acp.NewSessionResponse
	if err := json.Unmarshal(m.Result, &resp); err != nil {
		c.t.Fatalf("session/new result: %v", err)
	}
	if resp.SessionID == "" {
		c.t.Fatal("session/new returned no session id")
	}
	return resp.SessionID
}

func (c *testClient) prompt(id, text string) acp.PromptResponse {
	c.t.Helper()
	m := c.call(acp.MethodPrompt, acp.PromptRequest{
		SessionID: id,
		Prompt:    []acp.ContentBlock{{Type: "text", Text: text}},
	})
	if m.Error != nil {
		c.t.Fatalf("session/prompt: %v", m.Error)
	}
	var resp acp.PromptResponse
	if err := json.Unmarshal(m.Result, &resp); err != nil {
		c.t.Fatalf("session/prompt result: %v", err)
	}
	return resp
}

// --- the tests ------------------------------------------------------------

func TestOneTurnEndToEnd(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("the tests pass")))
	c := newTestServer(t, f, Options{})

	id := c.handshake(t.TempDir())
	resp := c.prompt(id, "does it build?")

	if resp.StopReason != acp.StopEndTurn {
		t.Errorf("stopReason = %q, want %q", resp.StopReason, acp.StopEndTurn)
	}
	if got := c.text(); got != "the tests pass" {
		t.Errorf("text = %q", got)
	}
	for _, u := range c.sessionUpdates(id) {
		if u.SessionUpdate == "" {
			t.Error("an update went out with no tag")
		}
	}
}

func TestInitializeAdvertisesWhatWeCanActuallyDo(t *testing.T) {
	c := newTestServer(t, newFakeAPI(t), Options{Version: "v9.9.9"})

	m := c.call(acp.MethodInitialize, acp.InitializeRequest{ProtocolVersion: 1})
	var resp acp.InitializeResponse
	if err := json.Unmarshal(m.Result, &resp); err != nil {
		t.Fatal(err)
	}

	if resp.AgentInfo == nil || resp.AgentInfo.Name != "codey" || resp.AgentInfo.Version != "v9.9.9" {
		t.Errorf("agentInfo = %+v", resp.AgentInfo)
	}
	if resp.AuthMethods == nil {
		t.Error("authMethods must be present, even empty")
	}
	// Session loading is not implemented, so it must not be advertised: a
	// client that believes it can resume will ask, and be refused.
	if resp.AgentCapabilities != nil && resp.AgentCapabilities.LoadSession {
		t.Error("loadSession is advertised but not implemented")
	}
}

func TestAuthenticateNeverFails(t *testing.T) {
	c := newTestServer(t, newFakeAPI(t), Options{})

	// Including for a method we never offered, which is what clients send.
	if m := c.call(acp.MethodAuthenticate, acp.AuthenticateRequest{MethodID: "xai.api_key"}); m.Error != nil {
		t.Fatalf("authenticate refused a method id: %v", m.Error)
	}
}

func TestAPromptForAnUnknownSessionIsAnError(t *testing.T) {
	c := newTestServer(t, newFakeAPI(t), Options{})
	c.handshake(t.TempDir())

	m := c.call(acp.MethodPrompt, acp.PromptRequest{SessionID: "nope", Prompt: []acp.ContentBlock{{Type: "text", Text: "hi"}}})
	if m.Error == nil {
		t.Fatal("want an error for an unknown session")
	}
	if m.Error.Code != acp.CodeInvalidParams {
		t.Errorf("code = %d", m.Error.Code)
	}
}

func TestCancelEndsTheTurnWithoutFailingIt(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("never sent")))
	f.hold = make(chan struct{})
	c := newTestServer(t, f, Options{})

	id := c.handshake(t.TempDir())

	done := make(chan acp.PromptResponse, 1)
	go func() { done <- c.prompt(id, "take your time") }()

	// The turn is now blocked inside the model call.
	time.Sleep(50 * time.Millisecond)
	c.notify(acp.MethodCancel, acp.CancelParams{SessionID: id})

	select {
	case resp := <-done:
		if resp.StopReason != acp.StopCancelled {
			t.Errorf("stopReason = %q, want cancelled", resp.StopReason)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("the cancel did not reach the turn")
	}
	close(f.hold)
}

func TestACancelBetweenTurnsDoesNotPoisonTheNextOne(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("one")), msg("end_turn", textBlock("two")))
	c := newTestServer(t, f, Options{})

	id := c.handshake(t.TempDir())
	c.prompt(id, "first")

	// Nothing is running; this must be forgotten rather than remembered.
	c.notify(acp.MethodCancel, acp.CancelParams{SessionID: id})
	time.Sleep(20 * time.Millisecond)

	if resp := c.prompt(id, "second"); resp.StopReason != acp.StopEndTurn {
		t.Errorf("stopReason = %q, want end_turn", resp.StopReason)
	}
}

func TestTwoSessionsDoNotCrossOver(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("alpha")), msg("end_turn", textBlock("beta")))
	c := newTestServer(t, f, Options{})

	c.handshake(t.TempDir()) // initialize once
	a := c.handshake(t.TempDir())
	b := c.handshake(t.TempDir())
	if a == b {
		t.Fatal("two sessions were given the same id")
	}

	c.prompt(a, "one")
	c.prompt(b, "two")

	for _, u := range c.sessionUpdates(a) {
		raw, _ := json.Marshal(u.Content)
		if strings.Contains(string(raw), "beta") {
			t.Error("session b's text arrived on session a")
		}
	}
}

func TestAToolRunsOnlyWhenThePermissionIsGranted(t *testing.T) {
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("t1", "write_file", `{"path":"note.txt","content":"hello"}`)),
		msg("end_turn", textBlock("written")),
	)
	c := newTestServer(t, f, Options{Permission: ModeAsk})

	root := t.TempDir()
	id := c.handshake(root)

	var asked acp.RequestPermissionParams
	c.onPermission = func(p acp.RequestPermissionParams) acp.RequestPermissionResult {
		asked = p
		return acp.RequestPermissionResult{
			Outcome: acp.PermissionOutcome{Outcome: acp.OutcomeSelected, OptionID: acp.OptAllowOnce},
		}
	}

	c.prompt(id, "write a note")

	if asked.SessionID != id {
		t.Fatalf("no permission was asked for, or for the wrong session: %+v", asked)
	}
	if asked.ToolCall.ToolCallID != "t1" {
		t.Errorf("the prompt was not attached to the call it is about: %q", asked.ToolCall.ToolCallID)
	}
	if len(asked.Options) != 3 {
		t.Errorf("want three answers to choose from, got %d", len(asked.Options))
	}
	if _, err := os.Stat(filepath.Join(root, "note.txt")); err != nil {
		t.Errorf("the file was not written after the permission was granted: %v", err)
	}
}

func TestARejectedToolDoesNotRun(t *testing.T) {
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("t1", "write_file", `{"path":"note.txt","content":"hello"}`)),
		msg("end_turn", textBlock("understood")),
	)
	c := newTestServer(t, f, Options{Permission: ModeAsk})

	root := t.TempDir()
	id := c.handshake(root)
	// The default answer is a rejection.
	c.prompt(id, "write a note")

	if _, err := os.Stat(filepath.Join(root, "note.txt")); !os.IsNotExist(err) {
		t.Error("the file was written despite the rejection")
	}
}

func TestFullAccessNeverAsks(t *testing.T) {
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("t1", "write_file", `{"path":"note.txt","content":"hello"}`)),
		msg("end_turn", textBlock("written")),
	)
	c := newTestServer(t, f, Options{Permission: ModeFullAccess})
	c.onPermission = func(acp.RequestPermissionParams) acp.RequestPermissionResult {
		t.Error("full access asked for permission")
		return acp.RequestPermissionResult{Outcome: acp.PermissionOutcome{Outcome: acp.OutcomeCancelled}}
	}

	root := t.TempDir()
	id := c.handshake(root)
	c.prompt(id, "write a note")

	if _, err := os.Stat(filepath.Join(root, "note.txt")); err != nil {
		t.Errorf("full access did not run the tool: %v", err)
	}
}

func TestTheModelPickerIsOfferedWhenThereIsOne(t *testing.T) {
	c := newTestServer(t, newFakeAPI(t), Options{
		Models:       []acp.ModelInfo{{ModelID: "a", Name: "A"}, {ModelID: "b", Name: "B"}},
		DefaultModel: "b",
	})

	c.call(acp.MethodInitialize, acp.InitializeRequest{ProtocolVersion: 1})
	m := c.call(acp.MethodNewSession, acp.NewSessionRequest{Cwd: t.TempDir()})

	var resp acp.NewSessionResponse
	if err := json.Unmarshal(m.Result, &resp); err != nil {
		t.Fatal(err)
	}
	if resp.Models == nil || len(resp.Models.AvailableModels) != 2 {
		t.Fatalf("models = %+v", resp.Models)
	}
	if resp.Models.CurrentModelID != "b" {
		t.Errorf("current model = %q", resp.Models.CurrentModelID)
	}
}

func TestSelectingAModelKeepsTheConversation(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("first")), msg("end_turn", textBlock("second")))

	var built []string
	c := newTestServer(t, f, Options{
		Models:       []acp.ModelInfo{{ModelID: "a", Name: "A"}, {ModelID: "b", Name: "B"}},
		DefaultModel: "a",
		NewAgent: func(root, model string, state *agent.State) (*agent.Agent, error) {
			built = append(built, model)
			log := metrics.NewWriter(io.Discard, "test")
			reg := tools.NewRegistry(log).Register(tools.NewRead(root, state))
			client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"), llm.WithModel(model))
			return agent.New(client, reg, state, log, agent.Options{}), nil
		},
	})

	id := c.handshake(t.TempDir())
	c.prompt(id, "first question")

	if m := c.call("session/set_model", acp.SetModelRequest{SessionID: id, ModelID: "b"}); m.Error != nil {
		t.Fatalf("session/set_model: %v", m.Error)
	}
	c.prompt(id, "second question")

	if len(built) != 2 || built[0] != "a" || built[1] != "b" {
		t.Fatalf("agents built for models %v, want [a b]", built)
	}
	// Switching models is not starting over: the second request must still
	// carry the first exchange.
	f.mu.Lock()
	defer f.mu.Unlock()
	if len(f.requests) != 2 {
		t.Fatalf("the fake API saw %d requests, want 2", len(f.requests))
	}
	second := string(f.requests[1])
	for _, want := range []string{"first question", "first", "second question"} {
		if !strings.Contains(second, want) {
			t.Errorf("the second request lost %q:\n%s", want, second)
		}
	}
	if !strings.Contains(second, `"model":"b"`) {
		t.Errorf("the second request did not go to the new model:\n%s", second)
	}
}

func TestAnUnknownConfigOptionIsAcceptedRatherThanRefused(t *testing.T) {
	c := newTestServer(t, newFakeAPI(t), Options{})
	c.call(acp.MethodInitialize, acp.InitializeRequest{ProtocolVersion: 1})

	m := c.call("session/set_config_option", map[string]any{"sessionId": "whatever", "configOptionId": "thinking", "value": "on"})
	if m.Error != nil {
		t.Fatalf("a config option we do not have should not fail a turn: %v", m.Error)
	}
}

func TestAnUnknownMethodIsRefusedCleanly(t *testing.T) {
	c := newTestServer(t, newFakeAPI(t), Options{})

	m := c.call("session/load", map[string]any{"sessionId": "x"})
	if m.Error == nil || m.Error.Code != acp.CodeMethodNotFound {
		t.Fatalf("want method not found, got %+v", m.Error)
	}
}

func TestImagesReachTheModel(t *testing.T) {
	// A 1x1 png, as a client would send it.
	const png = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="

	blocks := promptBlocks([]acp.ContentBlock{
		{Type: "text", Text: "what is this?"},
		{Type: "image", MimeType: "image/png", Data: png},
	})

	if len(blocks) != 2 {
		t.Fatalf("want a text block and an image block, got %d", len(blocks))
	}
	if blocks[0].Type != llm.BlockText || blocks[0].Text != "what is this?" {
		t.Errorf("first block = %+v", blocks[0])
	}
	if blocks[1].Type != llm.BlockImage || blocks[1].Source == nil || blocks[1].Source.Data != png {
		t.Errorf("the image did not survive: %+v", blocks[1])
	}
}

func TestAnAttachmentWeCannotDecodeIsStillMentioned(t *testing.T) {
	blocks := promptBlocks([]acp.ContentBlock{
		{Type: "text", Text: "look at this"},
		{Type: "resource_link", Name: "design.fig", URI: "file:///tmp/design.fig"},
	})

	if len(blocks) != 1 {
		t.Fatalf("want one text block, got %d", len(blocks))
	}
	if !strings.Contains(blocks[0].Text, "design.fig") {
		t.Errorf("the attachment was dropped silently: %q", blocks[0].Text)
	}
}