package serve

import (
	"context"
	"io"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/johannSo/codey/internal/acp"
	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/llm"
	"github.com/johannSo/codey/internal/metrics"
	"github.com/johannSo/codey/internal/subagent"
	"github.com/johannSo/codey/internal/tools"
)

// fakeAgents is a supervisor the test drives by hand: no children, no model
// calls, just the two callbacks and a snapshot the test controls.
type fakeAgents struct {
	mu       sync.Mutex
	observe  func(subagent.Event)
	changed  func()
	infos    []subagent.Info
	reports  []subagent.Info
	claimed  map[string]bool
	shutdown bool
}

func newFakeAgents() *fakeAgents { return &fakeAgents{claimed: map[string]bool{}} }

func (f *fakeAgents) Subscribe(observe func(subagent.Event), changed func()) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.observe, f.changed = observe, changed
}

func (f *fakeAgents) Snapshot() []subagent.Info {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]subagent.Info(nil), f.infos...)
}

func (f *fakeAgents) TakeReports() []subagent.Info {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := f.reports
	f.reports = nil
	return out
}

func (f *fakeAgents) ClaimReport(id string) bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.claimed[id] {
		return false
	}
	f.claimed[id] = true
	return true
}

func (f *fakeAgents) Shutdown() {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.shutdown = true
}

// say makes a child speak.
func (f *fakeAgents) say(id, preset, text string) {
	f.mu.Lock()
	observe := f.observe
	f.mu.Unlock()
	observe(subagent.Event{ID: id, Preset: preset, Model: "test-model",
		Agent: agent.Event{Kind: agent.EventText, Text: text}})
}

// finish ends a child, queues its report, and fires the change callback.
func (f *fakeAgents) finish(info subagent.Info) {
	f.mu.Lock()
	f.infos = append(f.infos, info)
	if info.Status == subagent.Done {
		f.reports = append(f.reports, info)
	}
	changed := f.changed
	f.mu.Unlock()
	changed()
}

// live adds a child that has not finished.
func (f *fakeAgents) live(info subagent.Info) {
	f.mu.Lock()
	f.infos = append(f.infos, info)
	f.mu.Unlock()
}

// --- the row translation, tested directly --------------------------------

func TestASubagentGetsItsOwnRow(t *testing.T) {
	e, r := newRecorded()
	sess := &session{emit: e, agentsMoved: make(chan struct{}, 1)}
	f := newFakeAgents()
	sess.agents = f
	sess.watchAgents()

	f.say("a1", "explore", strings.Repeat("looking around. ", 40))

	if len(r.updates) < 2 {
		t.Fatalf("want a row and an update, got %v", r.kinds())
	}
	open := r.updates[0]
	if open.SessionUpdate != acp.UpdateToolCall {
		t.Fatalf("first update = %q", open.SessionUpdate)
	}
	if open.ToolCallID != "agent-a1" {
		t.Errorf("toolCallId = %q — it must not collide with a real tool call", open.ToolCallID)
	}
	if !strings.Contains(open.Title, "a1") || !strings.Contains(open.Title, "explore") {
		t.Errorf("title = %q", open.Title)
	}
	if open.Status != acp.StatusInProgress {
		t.Errorf("status = %q", open.Status)
	}
	if !strings.Contains(contentText(t, r.updates[1]), "looking around") {
		t.Errorf("the child's words did not reach its row")
	}
}

func TestARowIsOpenedOnceAndClosedOnce(t *testing.T) {
	e, r := newRecorded()
	sess := &session{emit: e, agentsMoved: make(chan struct{}, 1)}
	f := newFakeAgents()
	sess.agents = f
	sess.watchAgents()

	f.say("a1", "explore", "one")
	r.clock = r.clock.Add(progressAfter)
	f.say("a1", "explore", "two")

	info := subagent.Info{ID: "a1", Preset: "explore", Status: subagent.Done, Turns: 3, Report: "found it"}
	f.finish(info)
	// The change callback fires repeatedly once a child has ended.
	f.finish(info)
	f.finish(info)

	var opens, closes int
	for _, u := range r.updates {
		switch {
		case u.SessionUpdate == acp.UpdateToolCall:
			opens++
		case u.Status == acp.StatusCompleted || u.Status == acp.StatusFailed:
			closes++
		}
	}
	if opens != 1 {
		t.Errorf("%d rows opened for one agent", opens)
	}
	if closes != 1 {
		t.Errorf("%d closing updates for one agent", closes)
	}
	if got := contentText(t, r.updates[len(r.updates)-1]); !strings.Contains(got, "found it") {
		t.Errorf("the row closed without the report: %q", got)
	}
}

func TestAFailedSubagentClosesAsFailed(t *testing.T) {
	e, r := newRecorded()
	sess := &session{emit: e, agentsMoved: make(chan struct{}, 1)}
	f := newFakeAgents()
	sess.agents = f
	sess.watchAgents()

	f.finish(subagent.Info{ID: "a1", Preset: "plan", Status: subagent.Failed, Err: "the endpoint refused"})

	last := r.updates[len(r.updates)-1]
	if last.Status != acp.StatusFailed {
		t.Errorf("status = %q, want failed", last.Status)
	}
	if !strings.Contains(contentText(t, last), "the endpoint refused") {
		t.Errorf("content = %q", contentText(t, last))
	}
	// It never said anything, so the row has to be opened before it is closed
	// or there is nothing on screen to close.
	if r.updates[0].SessionUpdate != acp.UpdateToolCall {
		t.Errorf("a silent agent left no row: %v", r.kinds())
	}
}

// --- the turn, over the wire ---------------------------------------------

// newAgentsServer wires a server whose sessions all share one fake supervisor.
func newAgentsServer(t *testing.T, f *fakeAPI, agents Agents) *testClient {
	t.Helper()
	log := metrics.NewWriter(io.Discard, "test")
	return newTestServer(t, f, Options{
		NewAgent: func(req AgentRequest) (AgentSet, error) {
			reg := tools.NewRegistry(log).Register(tools.NewRead(req.Root, req.State))
			client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"), llm.WithModel("test-model"))
			return AgentSet{
				Agent:  agent.New(client, reg, req.State, log, agent.Options{}),
				Agents: agents,
			}, nil
		},
	})
}

func TestATurnStaysOpenWhileASubagentWorks(t *testing.T) {
	// Ending the turn while a child is running would leave its row nowhere to
	// appear and its report nobody to be handed to.
	f := newFakeAPI(t,
		msg("end_turn", textBlock("spawned it")),
		msg("end_turn", textBlock("and here is what it found")),
	)
	agents := newFakeAgents()
	agents.live(subagent.Info{ID: "a1", Preset: "explore", Status: subagent.Running})

	c := newAgentsServer(t, f, agents)
	id := c.handshake(t.TempDir())

	done := make(chan acp.PromptResponse, 1)
	go func() { done <- c.prompt(id, "find the bug") }()

	select {
	case <-done:
		t.Fatal("the turn ended while a subagent was still running")
	case <-time.After(150 * time.Millisecond):
	}

	// The child finishes and reports. That restarts the turn with its report.
	agents.mu.Lock()
	agents.infos = nil
	agents.mu.Unlock()
	agents.finish(subagent.Info{ID: "a1", Preset: "explore", Status: subagent.Done, Turns: 2, Report: "it is the parser"})

	select {
	case resp := <-done:
		if resp.StopReason != acp.StopEndTurn {
			t.Errorf("stopReason = %q", resp.StopReason)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("the turn never ended")
	}

	// The report has to have reached the model, not only the screen.
	f.mu.Lock()
	defer f.mu.Unlock()
	if len(f.requests) != 2 {
		t.Fatalf("the parent was asked %d times, want 2", len(f.requests))
	}
	if !strings.Contains(string(f.requests[1]), "it is the parser") {
		t.Error("the subagent's report was never handed to the parent")
	}
}

func TestAReportAlreadyClaimedIsNotHandedOverTwice(t *testing.T) {
	// agent_wait may have collected it. Putting it in again would charge the
	// model twice for something already in its context.
	f := newFakeAPI(t, msg("end_turn", textBlock("done")))
	agents := newFakeAgents()
	agents.claimed["a1"] = true
	agents.reports = []subagent.Info{{ID: "a1", Preset: "explore", Status: subagent.Done, Report: "already seen"}}

	c := newAgentsServer(t, f, agents)
	id := c.handshake(t.TempDir())
	c.prompt(id, "go")

	f.mu.Lock()
	defer f.mu.Unlock()
	if len(f.requests) != 1 {
		t.Fatalf("the parent was asked %d times, want 1", len(f.requests))
	}
}

func TestCancellingATurnStopsWaitingForChildren(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("spawned")))
	agents := newFakeAgents()
	agents.live(subagent.Info{ID: "a1", Preset: "explore", Status: subagent.Running})

	c := newAgentsServer(t, f, agents)
	id := c.handshake(t.TempDir())

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

	time.Sleep(100 * 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("cancelling did not break the wait")
	}
}

func TestTheChildrenAreStoppedWhenTheConnectionEnds(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("ok")))
	agents := newFakeAgents()

	toAgent, fromClient := io.Pipe()
	fromAgent, toClient := io.Pipe()
	go func() { _, _ = io.Copy(io.Discard, fromAgent) }()

	log := metrics.NewWriter(io.Discard, "test")
	srv := New(Options{
		NewAgent: func(req AgentRequest) (AgentSet, error) {
			reg := tools.NewRegistry(log).Register(tools.NewRead(req.Root, req.State))
			client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"), llm.WithModel("m"))
			return AgentSet{Agent: agent.New(client, reg, req.State, log, agent.Options{}), Agents: agents}, nil
		},
	})

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

	// Open a session, then hang up.
	if _, err := io.WriteString(fromClient, `{"jsonrpc":"2.0","id":1,"method":"session/new","params":{"cwd":`+quote(t.TempDir())+"}}\n"); err != nil {
		t.Fatal(err)
	}
	time.Sleep(100 * time.Millisecond)
	_ = fromClient.Close()

	select {
	case <-served:
	case <-time.After(3 * time.Second):
		t.Fatal("Run did not return")
	}

	agents.mu.Lock()
	defer agents.mu.Unlock()
	if !agents.shutdown {
		t.Error("a subagent outlived the connection, still editing a tree nobody is watching")
	}
}

func TestSubagentRowsAreBoundedLikeAnyOtherToolOutput(t *testing.T) {
	e, r := newRecorded()
	sess := &session{emit: e, agentsMoved: make(chan struct{}, 1)}
	f := newFakeAgents()
	sess.agents = f
	sess.watchAgents()

	f.say("a1", "explore", strings.Repeat("x", progressMaxBytes*2)+"END")

	body := contentText(t, r.updates[len(r.updates)-1])
	if len(body) > progressMaxBytes+64 {
		t.Errorf("an unbounded %d bytes went over the wire", len(body))
	}
	// The child's line is newline-terminated, so trim before looking at the end.
	if !strings.HasSuffix(strings.TrimSpace(body), "END") {
		t.Error("the tail is the useful part and it was the part dropped")
	}
}