package subagent

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

	"codey/internal/agent"
	"codey/internal/approval"
	"codey/internal/metrics"
	"codey/internal/tools"
)

// Parallel agents are parallel spend, so only so many run at once and the rest
// wait their turn.
func TestTheConcurrencyCapQueuesTheRest(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f, func(o *Options) { o.MaxRunning = 1 })

	first, _ := s.Spawn(Spec{Preset: "explore", Task: "one"})
	second, _ := s.Spawn(Spec{Preset: "explore", Task: "two"})
	third, _ := s.Spawn(Spec{Preset: "explore", Task: "three"})

	waitFor(t, "the first agent to start", func() bool { return statusOf(t, s, first) == Running })
	time.Sleep(settle)
	if got := statusOf(t, s, second); got != Queued {
		t.Errorf("second agent is %s, want queued behind the cap", got)
	}
	if got := statusOf(t, s, third); got != Queued {
		t.Errorf("third agent is %s, want queued behind the cap", got)
	}
	if n := f.count(); n != 1 {
		t.Errorf("%d agents were talking to the API at once, want 1", n)
	}

	f.release()
	waitFor(t, "every agent to finish", func() bool {
		return statusOf(t, s, first) == Done && statusOf(t, s, second) == Done && statusOf(t, s, third) == Done
	})
}

func TestKillStopsARunningAgent(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f)

	id, _ := s.Spawn(Spec{Preset: "general", Task: "go"})
	waitFor(t, "the agent to start", func() bool { return statusOf(t, s, id) == Running })

	if err := s.Kill(id); err != nil {
		t.Fatalf("Kill: %v", err)
	}
	waitFor(t, "the agent to stop", func() bool { return statusOf(t, s, id) == Killed })
	f.release()

	info, _ := s.Info(id)
	if info.Err != "" {
		t.Errorf("a killed agent recorded an error (%q); being stopped is not a failure", info.Err)
	}
	if !strings.Contains(info.Notification(), "stopped by the user") {
		t.Errorf("the parent would be told: %q", info.Notification())
	}
}

// Killing something that has not started must not leave it queued, waiting to
// start after it was stopped.
func TestKillingAQueuedAgentNeverRunsIt(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f, func(o *Options) { o.MaxRunning = 1 })

	first, _ := s.Spawn(Spec{Preset: "explore", Task: "one"})
	queued, _ := s.Spawn(Spec{Preset: "explore", Task: "two"})
	waitFor(t, "the first agent to start", func() bool { return statusOf(t, s, first) == Running })

	if err := s.Kill(queued); err != nil {
		t.Fatalf("Kill: %v", err)
	}
	if got := statusOf(t, s, queued); got != Killed {
		t.Fatalf("queued agent is %s after being killed", got)
	}

	f.release()
	waitFor(t, "the first agent to finish", func() bool { return statusOf(t, s, first) == Done })
	time.Sleep(settle)
	if n := f.count(); n != 1 {
		t.Errorf("the API saw %d requests, want 1 — the killed agent still ran", n)
	}
}

// Pausing holds the agent at the next turn boundary and spends nothing until
// it is resumed.
func TestPauseHoldsAnAgentAndResumeContinuesIt(t *testing.T) {
	f := newFakeAPI(t)
	f.twoTurns = true
	var sup *Supervisor
	// Pausing from inside the first turn's tool call is the only way to be
	// sure the request for the second turn has not gone out yet.
	s := newSupervisor(t, f, func(o *Options) {
		o.Observe = func(ev Event) {
			if ev.Agent.Kind == agent.EventToolResult {
				sup.Pause(ev.ID)
			}
		}
	})
	sup = s

	id, _ := s.Spawn(Spec{Preset: "general", Task: "go"})
	waitFor(t, "the agent to be held", func() bool { return statusOf(t, s, id) == Paused })
	time.Sleep(settle)
	if n := f.count(); n != 1 {
		t.Fatalf("a paused agent made %d requests, want 1", n)
	}

	if err := s.Resume(id); err != nil {
		t.Fatalf("Resume: %v", err)
	}
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })
	if n := f.count(); n != 2 {
		t.Errorf("made %d requests overall, want 2", n)
	}
}

// Pausing something that never started would be a lie in the sidebar: there is
// nothing to hold. Say so, and point at the thing that does work.
func TestPausingAQueuedAgentSaysItHasNotStarted(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f, func(o *Options) { o.MaxRunning = 1 })

	first, _ := s.Spawn(Spec{Preset: "explore", Task: "one"})
	queued, _ := s.Spawn(Spec{Preset: "explore", Task: "two"})
	waitFor(t, "the first agent to start", func() bool { return statusOf(t, s, first) == Running })

	err := s.Pause(queued)
	if err == nil {
		t.Fatal("pausing a queued agent succeeded")
	}
	if !strings.Contains(err.Error(), "not started") {
		t.Errorf("error = %q, want it to say the agent has not started", err)
	}
	if got := statusOf(t, s, queued); got != Queued {
		t.Errorf("a failed pause changed the status to %s", got)
	}
	f.release()
}

func TestPauseAndResumeOnAFinishedAgentAreRefused(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "go"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	for _, err := range []error{s.Pause(id), s.Resume(id), s.Kill(id)} {
		if err == nil {
			t.Error("controlling a finished agent succeeded")
		} else if !strings.Contains(err.Error(), "already finished") {
			t.Errorf("error = %q, want it to say the agent has finished", err)
		}
	}
}

func TestWaitReturnsHowTheAgentEnded(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "find it"})

	info, err := s.Wait(context.Background(), id)
	if err != nil {
		t.Fatalf("Wait: %v", err)
	}
	if info.Status != Done || info.Report != "handled: find it" {
		t.Errorf("Wait returned %+v", info)
	}
}

// Interrupting the turn that is waiting must not kill the agent it waited for:
// the wait is the parent's, the work is the child's.
func TestWaitIsInterruptibleWithoutKillingTheAgent(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "general", Task: "go"})

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := s.Wait(ctx, id); err == nil {
		t.Fatal("a cancelled Wait returned no error")
	}
	if got := statusOf(t, s, id); !got.Live() {
		t.Errorf("the agent is %s after its waiter gave up", got)
	}

	f.release()
	waitFor(t, "the agent to finish on its own", func() bool { return statusOf(t, s, id) == Done })
}

func TestWaitingOnAnUnknownAgentSaysSo(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))
	if _, err := s.Wait(context.Background(), "a9"); err == nil {
		t.Fatal("waiting on an agent that does not exist succeeded")
	}
}

// Reports are handed to the parent exactly once: twice would repeat the work
// in its context, never would lose it.
func TestReportsAreHandedOverOnce(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "go"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	first := s.TakeReports()
	if len(first) != 1 || first[0].ID != id {
		t.Fatalf("TakeReports = %+v, want the finished agent", first)
	}
	if again := s.TakeReports(); len(again) != 0 {
		t.Errorf("TakeReports handed the same report over twice: %+v", again)
	}
}

// The front end redraws on a callback rather than by polling, so a status
// change that fires nothing would leave a finished agent showing as running.
func TestChangedFiresOnStatusChanges(t *testing.T) {
	f := newFakeAPI(t)
	var changes int
	var mu = make(chan struct{}, 1)
	mu <- struct{}{}
	s := newSupervisor(t, f, func(o *Options) {
		o.Changed = func() {
			<-mu
			changes++
			mu <- struct{}{}
		}
	})

	id, _ := s.Spawn(Spec{Preset: "explore", Task: "go"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })
	waitFor(t, "the change callback", func() bool {
		<-mu
		defer func() { mu <- struct{}{} }()
		return changes >= 2
	})
}

// Leaving the harness must not leave agents editing a tree nobody is watching.
func TestShutdownStopsEverything(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f, func(o *Options) { o.MaxRunning = 1 })

	running, _ := s.Spawn(Spec{Preset: "general", Task: "one"})
	queued, _ := s.Spawn(Spec{Preset: "general", Task: "two"})
	waitFor(t, "the first agent to start", func() bool { return statusOf(t, s, running) == Running })

	done := make(chan struct{})
	go func() { s.Shutdown(); close(done) }()
	f.release()

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("Shutdown never returned")
	}
	for _, id := range []string{running, queued} {
		if got := statusOf(t, s, id); got.Live() {
			t.Errorf("agent %s is %s after shutdown", id, got)
		}
	}
}

// A subagent's tool calls go through the same policy the session's own do, and
// the request has to say which agent is asking — you cannot decide on a shell
// command without knowing who wants to run it. While it waits, the sidebar has
// to show that it is stopped on you rather than working.
func TestASubagentsApprovalIsTaggedAndShowsAsWaiting(t *testing.T) {
	f := newFakeAPI(t)
	f.callTool = "noop"

	asked := make(chan approval.Request, 4)
	release := make(chan approval.Decision)
	policy := approval.New(func(ctx context.Context, req approval.Request) (approval.Decision, error) {
		asked <- req
		select {
		case d := <-release:
			return d, nil
		case <-ctx.Done():
			return approval.Deny, ctx.Err()
		}
	})
	s := newSupervisor(t, f, func(o *Options) { o.Policy = policy })

	id, _ := s.Spawn(Spec{Preset: "general", Task: "go"})

	var req approval.Request
	select {
	case req = <-asked:
	case <-time.After(3 * time.Second):
		t.Fatal("the subagent's tool call never reached the approval policy")
	}
	if req.Agent != id {
		t.Errorf("the prompt says agent %q, want %q", req.Agent, id)
	}
	waitFor(t, "the agent to show as waiting", func() bool { return statusOf(t, s, id) == Waiting })

	release <- approval.Once
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })
}

// newBlockedSupervisor holds every child inside its tool-set construction, so
// a test can act on an agent in the window between being given a slot and
// having anything to cancel. release lets them through.
func newBlockedSupervisor(t *testing.T, f *fakeAPI) (*Supervisor, func()) {
	t.Helper()
	gate := make(chan struct{})
	root := t.TempDir()
	s := newSupervisor(t, f, func(o *Options) {
		o.Tools = func(state *agent.State, log *metrics.Logger) *tools.Registry {
			<-gate
			return tools.NewRegistry(log).Register(tools.NewRead(root, state))
		}
	})
	var once sync.Once
	return s, func() { once.Do(func() { close(gate) }) }
}

// Killing an agent that has a slot but has not reached its first request must
// stop it. Losing the kill there means watching an agent you stopped carry on
// editing files.
func TestKillingAnAgentThatIsStillStartingStopsIt(t *testing.T) {
	f := newFakeAPI(t)
	s, release := newBlockedSupervisor(t, f)
	defer release()

	// Spawn blocks in the gate: a child is built on the goroutine that asked
	// for it, and the gate is inside that construction.
	go s.Spawn(Spec{Preset: "general", Task: "go"})
	waitFor(t, "the agent to take a slot", func() bool {
		rows := s.Snapshot()
		return len(rows) == 1 && rows[0].Status == Running
	})
	id := s.Snapshot()[0].ID

	if err := s.Kill(id); err != nil {
		t.Fatalf("Kill: %v", err)
	}
	release()
	waitFor(t, "the agent to stop", func() bool { return statusOf(t, s, id) == Killed })

	time.Sleep(settle)
	if n := f.count(); n != 0 {
		t.Errorf("a killed agent still made %d requests", n)
	}
}

// Shutting down in that same window must not find a half-built child it
// accounts for twice — once as never-started and once when it ends.
func TestShutdownDuringStartupDoesNotPanic(t *testing.T) {
	f := newFakeAPI(t)
	s, release := newBlockedSupervisor(t, f)

	go s.Spawn(Spec{Preset: "general", Task: "go"})
	waitFor(t, "the agent to take a slot", func() bool {
		rows := s.Snapshot()
		return len(rows) == 1 && rows[0].Status == Running
	})
	id := s.Snapshot()[0].ID

	done := make(chan struct{})
	go func() { s.Shutdown(); close(done) }()
	time.Sleep(settle)
	release()

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("Shutdown never returned")
	}
	if got := statusOf(t, s, id); got.Live() {
		t.Errorf("agent is %s after shutdown", got)
	}

	// The half-built child is let go last. If shutting down accounted for it
	// once and its own ending accounts for it again, this is where the process
	// dies — so the test waits for that ending rather than leaving it to land
	// after the run, as a crash somebody sees once a month.
	waitFor(t, "the released agent to end", func() bool { return statusOf(t, s, id) == Killed })
	time.Sleep(settle)
	if n := f.count(); n != 0 {
		t.Errorf("an agent shut down before it started still made %d requests", n)
	}
}

// The sequence that used to send every report twice: the agent finishes, the
// front end drains it into its pending queue, and only then does the parent's
// agent_wait land. Whoever hands the report to the model has to claim it, or
// the other one hands it over again — which is what the whole feature exists
// to avoid.
func TestAReportDrainedBeforeTheWaitIsNotSentTwice(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "find it"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	// The front end drains first — it holds the report, it has not given it
	// to the model yet.
	if reports := s.TakeReports(); len(reports) != 1 {
		t.Fatalf("TakeReports = %+v", reports)
	}

	// Now the parent asks for it, and gets it as the tool's result.
	info, err := s.Wait(context.Background(), id)
	if err != nil {
		t.Fatalf("Wait: %v", err)
	}
	if !s.ClaimReport(info.ID) {
		t.Fatal("agent_wait could not claim a report nobody had delivered")
	}

	// The front end, coming to inject the copy it drained, must be refused.
	if s.ClaimReport(id) {
		t.Error("the front end would put the same report into the conversation a second time")
	}
}

// Claiming is the handover, so exactly one caller can win it however many ask.
func TestOnlyOneClaimOnAReportEverWins(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "find it"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	var won int
	for i := 0; i < 4; i++ {
		if s.ClaimReport(id) {
			won++
		}
	}
	if won != 1 {
		t.Errorf("%d callers were told to deliver the same report", won)
	}
	if s.ClaimReport("a99") {
		t.Error("an agent that does not exist had a report to claim")
	}
}

// An agent still working has nothing to hand over yet.
func TestARunningAgentsReportCannotBeClaimed(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "general", Task: "go"})
	waitFor(t, "the agent to start", func() bool { return statusOf(t, s, id) == Running })

	if s.ClaimReport(id) {
		t.Error("claimed a report from an agent that has not finished")
	}
	f.release()
}

// A wait the user interrupted delivered nothing, so the report must still be
// there for the front end to hand over.
func TestAnInterruptedWaitLeavesTheReportToBeDelivered(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "find it"})

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := s.Wait(ctx, id); err == nil {
		t.Fatal("a cancelled Wait returned no error")
	}

	f.release()
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })
	if !s.ClaimReport(id) {
		t.Error("the report was claimed by a wait that was interrupted before it delivered anything")
	}
}