// Package subagent runs child agents alongside the session's own.
//
// A subagent is a whole agent: its own conversation, its own context window,
// its own model, its own goroutine. The agent that spawned one never sees its
// turns — only the report it ends with — which is the point. Work that would
// have filled the parent's window is spent somewhere else and comes back a
// paragraph long.
//
// Like the agent package, this one knows nothing about a terminal. It reports
// through callbacks and a snapshot, so the TUI and the eval runner can both
// drive it.
package subagent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

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

// Status is where one subagent is in its life.
type Status string

const (
	// Queued is spawned but waiting for a slot: only so many agents run at
	// once, because parallel agents are parallel spend.
	Queued Status = "queued"
	// Waiting is blocked on an approval prompt — stopped by you, not by
	// itself.
	Waiting Status = "waiting"
	Running Status = "running"
	Paused  Status = "paused"
	Done    Status = "done"
	Failed  Status = "failed"
	Killed  Status = "killed"
)

// Live reports whether the agent is still going to do something. A live agent
// holds one of the concurrency slots; a finished one does not.
func (s Status) Live() bool {
	switch s {
	case Queued, Running, Waiting, Paused:
		return true
	}
	return false
}

// DefaultMaxRunning caps how many subagents run at once. Four is enough to
// parallelise a search and few enough that a runaway spawn loop is visible
// before it is expensive.
const DefaultMaxRunning = 4

// Info is one subagent as the outside world sees it: everything the sidebar
// draws and everything the parent is told.
type Info struct {
	ID     string
	Preset string
	Model  string
	// Effort is the level this agent runs at, empty when it runs at the
	// session's — or when the model it is on takes no effort level at all.
	Effort string
	Task   string
	Status Status
	Turns  int
	Usage  llm.Usage
	// Started is when it got a slot, not when it was spawned; a queued agent
	// has not started.
	Started time.Time
	Ended   time.Time
	// Report is the agent's final message: the whole of what it hands back.
	Report string
	// Err is why it failed, empty otherwise.
	Err string
}

// Cost prices the agent's usage at its own model, which is the only honest way
// to bill a haiku agent spawned from an opus session.
func (i Info) Cost() (float64, bool) { return i.Usage.Cost(i.Model) }

// Elapsed is how long it has been running, or how long it ran.
func (i Info) Elapsed() time.Duration {
	if i.Started.IsZero() {
		return 0
	}
	if i.Ended.IsZero() {
		return time.Since(i.Started)
	}
	return i.Ended.Sub(i.Started)
}

// Notification is what the parent agent is told when this agent finishes. It
// goes into the conversation as a user message, so it names the agent and
// repeats the task: by the time it lands, several turns may have passed and
// the parent has to be able to place it.
func (i Info) Notification() string {
	var b strings.Builder
	switch i.Status {
	case Done:
		fmt.Fprintf(&b, "[subagent %s (%s) finished after %s]\n", i.ID, i.Preset, plural(i.Turns, "turn"))
		fmt.Fprintf(&b, "Its task was: %s\n\nIts report:\n%s", i.Task, i.Report)
	case Failed:
		fmt.Fprintf(&b, "[subagent %s (%s) failed]\nIts task was: %s\nIt stopped because: %s\n\n", i.ID, i.Preset, i.Task, i.Err)
		b.WriteString("Nothing it may have done was rolled back. Decide whether to do this work yourself, spawn another agent, or leave it.")
		if i.Report != "" {
			fmt.Fprintf(&b, "\n\nWhat it had said before it stopped:\n%s", i.Report)
		}
	case Killed:
		fmt.Fprintf(&b, "[subagent %s (%s) was stopped by the user]\nIts task was: %s\n\n", i.ID, i.Preset, i.Task)
		b.WriteString("Do not spawn it again unless you are asked to: the user stopped it on purpose.")
	default:
		fmt.Fprintf(&b, "[subagent %s (%s) is %s]", i.ID, i.Preset, i.Status)
	}
	return b.String()
}

// AlreadyDelivered is what agent_wait answers when the report has already
// reached the parent another way. Repeating the text would charge it twice for
// something already in its context, which is the cost this whole feature
// exists to avoid.
func (i Info) AlreadyDelivered() string {
	return fmt.Sprintf("[subagent %s (%s) %s after %s — its report has already been delivered to you "+
		"as a message in this conversation. Read it there rather than waiting again.]",
		i.ID, i.Preset, endedWord(i.Status), plural(i.Turns, "turn"))
}

// endedWord is how an agent's ending is described in a sentence.
func endedWord(s Status) string {
	switch s {
	case Failed:
		return "failed"
	case Killed:
		return "was stopped"
	}
	return "finished"
}

func plural(n int, word string) string {
	if n == 1 {
		return fmt.Sprintf("%d %s", n, word)
	}
	return fmt.Sprintf("%d %ss", n, word)
}

// Event is one child's agent event, tagged with the child it came from. The
// front end keeps a transcript per agent and needs to know which.
type Event struct {
	ID     string
	Preset string
	Model  string
	Agent  agent.Event
}

// Options wires a Supervisor.
type Options struct {
	// Client is the session's. Children are cloned off it so they
	// authenticate identically and differ only in the model.
	Client *llm.Client
	Log    *metrics.Logger
	// Policy gates children's tool calls. Nil means nothing is gated, which
	// is only right where the blast radius is already contained.
	Policy *approval.Policy
	// Base carries the session's agent settings — streaming, caching, the
	// context window, the identity prefix, the discovered skills. The preset
	// supplies the system prompt, the effort and the turn cap over the top.
	Base agent.Options
	// Tools builds the tool set for one child. It is supplied by the caller
	// because the tools need a root, and the supervisor has no business
	// knowing one. Note what it does not include: the spawn tools, so a
	// subagent cannot spawn subagents.
	Tools func(*agent.State, *metrics.Logger) *tools.Registry
	// MaxRunning caps concurrency. Zero means DefaultMaxRunning.
	MaxRunning int
	// Observe receives every child's events, tagged. It may be nil.
	Observe func(Event)
	// Changed is called when a status changes or a report is queued, so a
	// front end knows to redraw without polling. It may be nil.
	Changed func()
}

// child is one subagent and the handles needed to control it.
type child struct {
	info   Info
	ag     *agent.Agent
	cancel context.CancelFunc
	// killed records that the cancellation was deliberate, so the ordinary
	// context error is reported as a kill rather than as a failure.
	killed bool
	// pauseWanted holds a pause asked for before the agent existed.
	pauseWanted bool
	// delivered records that this agent's report has been handed to the
	// parent — as an agent_wait result or as a message the front end put into
	// the conversation. Whoever hands it over claims it first, so the other
	// one knows not to.
	delivered bool
	done      chan struct{}
}

// Supervisor owns every subagent of one session.
type Supervisor struct {
	ctx  context.Context
	opts Options

	mu    sync.Mutex
	byID  map[string]*child
	order []*child
	n     int
	// reports are finished agents the parent has not been told about yet.
	reports []Info

	wg sync.WaitGroup
}

// New builds a supervisor. ctx is the parent of every child's context: it
// outlives any one turn, because a subagent must survive the interruption of
// the turn that spawned it.
func New(ctx context.Context, opts Options) *Supervisor {
	if opts.MaxRunning <= 0 {
		opts.MaxRunning = DefaultMaxRunning
	}
	return &Supervisor{ctx: ctx, opts: opts, byID: map[string]*child{}}
}

// Subscribe sets the callbacks a front end wants. It exists because the TUI
// cannot be built before the supervisor it watches: the program has to exist
// to be sent messages, and the supervisor has to exist to be handed to it.
func (s *Supervisor) Subscribe(observe func(Event), changed func()) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.opts.Observe, s.opts.Changed = observe, changed
}

// observer and notifier read the callbacks under the lock, so Subscribe can be
// called while children are already running.
func (s *Supervisor) observer() func(Event) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.opts.Observe
}

func (s *Supervisor) changed() {
	s.mu.Lock()
	fn := s.opts.Changed
	s.mu.Unlock()
	if fn != nil {
		fn()
	}
}

// Spec is what one subagent is asked to be: which preset, what to do, and what
// to spend on it. Model and Effort are the caller's choice — the presets pin
// neither — and both may be empty, which inherits the session's.
type Spec struct {
	Preset string
	Task   string
	Model  string
	Effort string
}

// Spawn starts a subagent, or queues it when the concurrency cap is full. It
// returns the id immediately: the work happens on its own goroutine.
func (s *Supervisor) Spawn(spec Spec) (string, error) {
	p, err := Lookup(spec.Preset)
	if err != nil {
		return "", err
	}
	if strings.TrimSpace(spec.Task) == "" {
		return "", errors.New("a subagent needs a task: say what it should do and what to report back")
	}
	task := spec.Task

	model := spec.Model
	if model == "" {
		// Nobody chose: the session's own model is the one answer that is
		// certainly served by this endpoint.
		model = s.opts.Client.Model()
	}

	effort := spec.Effort
	if effort != "" && !llm.ValidEffort(effort) {
		return "", fmt.Errorf("no such effort level %q. The ones that exist are: %s",
			effort, strings.Join(llm.EffortLevels, ", "))
	}
	// A level this model does not take is not a reason to refuse the spawn —
	// the job and the model may both be right — but it is not in force, and
	// recording it as though it were would be a lie in the sidebar and in the
	// answer the parent gets.
	if !llm.SupportsEffort(model, effort) {
		effort = ""
	}

	s.mu.Lock()
	s.n++
	id := fmt.Sprintf("a%d", s.n)
	c := &child{
		info: Info{ID: id, Preset: p.Name, Model: model, Effort: effort, Task: task, Status: Queued},
		done: make(chan struct{}),
	}
	s.byID[id] = c
	s.order = append(s.order, c)
	s.mu.Unlock()

	s.opts.Log.Log(metrics.Event{Type: "subagent_spawn", Extra: map[string]any{
		"agent": id, "preset": p.Name, "model": model, "effort": effort,
	}})
	s.changed()
	s.pump()
	return id, nil
}

// pump starts queued children while there are slots. It is called whenever the
// set of running agents can have changed.
func (s *Supervisor) pump() {
	for {
		s.mu.Lock()
		running := 0
		for _, c := range s.order {
			if c.info.Status.Live() && c.info.Status != Queued {
				running++
			}
		}
		var next *child
		if running < s.opts.MaxRunning {
			for _, c := range s.order {
				if c.info.Status == Queued {
					next = c
					break
				}
			}
		}
		if next == nil {
			s.mu.Unlock()
			return
		}
		next.info.Status = Running
		next.info.Started = time.Now()
		// The context, its cancel, and the wait-group count are all taken here,
		// in the same hold of the lock that marks the child running. Setting
		// them up inside start instead would leave a window where the child
		// is running with nothing to cancel and nothing to wait for: a kill
		// landing in it would be lost, and a shutdown would account for the
		// child twice.
		ctx, cancel := context.WithCancel(s.ctx)
		next.cancel = cancel
		s.wg.Add(1)
		s.mu.Unlock()

		s.start(next, ctx, cancel)
		s.changed()
	}
}

// start builds the child's agent and runs it. The context is already live and
// already cancellable by the time it gets here.
func (s *Supervisor) start(c *child, ctx context.Context, cancel context.CancelFunc) {
	p, _ := Lookup(c.info.Preset)
	id, task := c.info.ID, c.info.Task

	log := s.opts.Log.For(id)
	state := agent.NewState()
	reg := s.opts.Tools(state, log)

	o := s.opts.Base
	o.System = p.System
	if c.info.Effort != "" {
		o.Effort = c.info.Effort
	}
	if p.MaxTurns > 0 {
		o.MaxTurns = p.MaxTurns
	}
	ag := agent.New(s.opts.Client.Clone(c.info.Model), reg, state, log, o)
	ag.Observe = func(ev agent.Event) { s.onEvent(id, ev) }
	if s.opts.Policy != nil {
		ag.Approve = approver{sup: s, id: id}
	}

	s.mu.Lock()
	c.ag = ag
	pause := c.pauseWanted
	s.mu.Unlock()
	if pause {
		ag.Pause()
	}

	go func() {
		defer s.wg.Done()
		defer cancel()
		// A kill that landed while this child was being built has already
		// cancelled ctx, so Run returns at once and finish records it as
		// killed — without the agent having done anything.
		report, err := ag.Run(ctx, task)
		s.finish(c, report, err)
	}()
}

// onEvent keeps the snapshot current and forwards the event.
func (s *Supervisor) onEvent(id string, ev agent.Event) {
	var status bool

	s.mu.Lock()
	c := s.byID[id]
	if c != nil {
		switch ev.Kind {
		case agent.EventTurnStart:
			c.info.Turns = ev.Turn
		case agent.EventUsage:
			c.info.Usage = addUsage(c.info.Usage, ev.Usage)
		case agent.EventPaused:
			c.info.Status, status = Paused, true
		case agent.EventResumed:
			c.info.Status, status = Running, true
		}
	}
	var preset, model string
	if c != nil {
		preset, model = c.info.Preset, c.info.Model
	}
	s.mu.Unlock()

	if fn := s.observer(); fn != nil {
		fn(Event{ID: id, Preset: preset, Model: model, Agent: ev})
	}
	if status {
		s.changed()
	}
}

// finish records how a child ended and queues its report for the parent.
func (s *Supervisor) finish(c *child, report string, err error) {
	s.mu.Lock()
	c.info.Ended = time.Now()
	c.info.Report = strings.TrimSpace(report)
	switch {
	case c.killed:
		c.info.Status = Killed
	case err != nil:
		c.info.Status = Failed
		c.info.Err = err.Error()
	default:
		c.info.Status = Done
	}
	info := c.info
	s.reports = append(s.reports, info)
	close(c.done)
	s.mu.Unlock()

	cost, _ := info.Cost()
	s.opts.Log.Log(metrics.Event{Type: "subagent_finish", Turn: info.Turns, Extra: map[string]any{
		"agent": info.ID, "preset": info.Preset, "status": string(info.Status),
		"seconds": info.Elapsed().Seconds(), "cost_usd": cost,
	}})
	s.changed()
	s.pump()
}

// Snapshot is every subagent this session has spawned, in spawn order.
func (s *Supervisor) Snapshot() []Info {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]Info, 0, len(s.order))
	for _, c := range s.order {
		out = append(out, c.info)
	}
	return out
}

// Info is one subagent by id.
func (s *Supervisor) Info(id string) (Info, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	c, ok := s.byID[id]
	if !ok {
		return Info{}, false
	}
	return c.info, true
}

// Live is how many subagents are still going, for the status line.
func (s *Supervisor) Live() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	var n int
	for _, c := range s.order {
		if c.info.Status.Live() {
			n++
		}
	}
	return n
}

// TakeReports hands over the finished agents the parent has not been told
// about, and forgets them. The caller is expected to deliver every one.
func (s *Supervisor) TakeReports() []Info {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := s.reports
	s.reports = nil
	return out
}

func (s *Supervisor) lookupLive(id, verb string) (*child, error) {
	c, ok := s.byID[id]
	if !ok {
		return nil, fmt.Errorf("no subagent %q. Check the id you were given when you spawned it", id)
	}
	if !c.info.Status.Live() {
		return nil, fmt.Errorf("subagent %s has already finished (%s), so it cannot be %s", id, c.info.Status, verb)
	}
	return c, nil
}

// Pause holds a subagent at its next turn boundary.
func (s *Supervisor) Pause(id string) error {
	s.mu.Lock()
	c, err := s.lookupLive(id, "paused")
	if err != nil {
		s.mu.Unlock()
		return err
	}
	if c.info.Status == Paused {
		s.mu.Unlock()
		return nil
	}
	if c.info.Status == Queued {
		// There is nothing to hold, and saying "paused" about an agent that
		// was never running would be a lie in the sidebar.
		s.mu.Unlock()
		return fmt.Errorf("subagent %s has not started yet — it is queued behind the ones running. Stop it if you do not want it to run", id)
	}
	// Marked before the turn in flight has finished, so the sidebar says what
	// was asked for rather than lagging a whole turn behind it. pauseWanted
	// covers the moment between getting a slot and having an agent to hold.
	c.pauseWanted = true
	c.info.Status = Paused
	ag := c.ag
	s.mu.Unlock()

	if ag != nil {
		ag.Pause()
	}
	s.changed()
	return nil
}

// Resume lets a paused subagent carry on.
func (s *Supervisor) Resume(id string) error {
	s.mu.Lock()
	c, err := s.lookupLive(id, "resumed")
	if err != nil {
		s.mu.Unlock()
		return err
	}
	c.pauseWanted = false
	if c.info.Status == Paused {
		c.info.Status = Running
	}
	ag := c.ag
	s.mu.Unlock()

	if ag != nil {
		ag.Resume()
	}
	s.changed()
	return nil
}

// Kill stops a subagent now. Whatever it has already done to the working tree
// stays done — this ends the agent, it does not undo it.
func (s *Supervisor) Kill(id string) error {
	s.mu.Lock()
	c, err := s.lookupLive(id, "stopped")
	if err != nil {
		s.mu.Unlock()
		return err
	}
	c.killed = true
	cancel := c.cancel
	neverStarted := c.info.Status == Queued
	if neverStarted {
		// Nothing is running to notice the cancellation, so this is where it
		// ends: recorded, reported, and out of the queue.
		c.info.Status = Killed
		c.info.Ended = time.Now()
		s.reports = append(s.reports, c.info)
		close(c.done)
	}
	s.mu.Unlock()

	if cancel != nil {
		cancel()
	}
	s.opts.Log.Log(metrics.Event{Type: "subagent_killed", Extra: map[string]any{"agent": id}})
	s.changed()
	if neverStarted {
		s.pump()
	}
	return nil
}

// Wait blocks until a subagent finishes and returns how it ended. The context
// is the caller's — interrupting the turn that is waiting does not kill the
// agent it was waiting for.
func (s *Supervisor) Wait(ctx context.Context, id string) (Info, error) {
	s.mu.Lock()
	c, ok := s.byID[id]
	s.mu.Unlock()
	if !ok {
		return Info{}, fmt.Errorf("no subagent %q. Check the id you were given when you spawned it", id)
	}

	select {
	case <-c.done:
	case <-ctx.Done():
		return Info{}, ctx.Err()
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	return c.info, nil
}

// ClaimReport asks for the right to give one finished agent's report to the
// parent, and reports whether the caller got it. Exactly one caller ever does.
//
// Two of them compete, and which one arrives first is a race nobody can
// arrange: agent_wait, which returns the report as its tool result, and the
// front end, which puts it into the conversation as a message. Claiming is the
// handover itself rather than a flag either side reads from a copy — a copy
// taken before the other side acted is exactly how the same report used to be
// handed over twice.
func (s *Supervisor) ClaimReport(id string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	c, ok := s.byID[id]
	// An agent still working has nothing to hand over yet.
	if !ok || c.info.Status.Live() || c.delivered {
		return false
	}
	c.delivered = true
	return true
}

// Shutdown kills every subagent and waits for them to unwind. It is what
// leaving Codey does: an agent outliving the session it belongs to would
// keep editing a tree nobody is watching.
func (s *Supervisor) Shutdown() {
	s.mu.Lock()
	var cancels []context.CancelFunc
	for _, c := range s.order {
		if !c.info.Status.Live() {
			continue
		}
		c.killed = true
		if c.cancel != nil {
			cancels = append(cancels, c.cancel)
			continue
		}
		// Queued and never started: nothing will report it, so record it here.
		c.info.Status = Killed
		c.info.Ended = time.Now()
		close(c.done)
	}
	s.mu.Unlock()

	for _, cancel := range cancels {
		cancel()
	}
	s.wg.Wait()
}

// approver routes a child's approval requests through the session's policy,
// tagged with the agent asking, and marks the agent as waiting while a human
// decides. Several agents can be waiting at once; the UI queues them.
type approver struct {
	sup *Supervisor
	id  string
}

func (a approver) Approve(ctx context.Context, tool string, input json.RawMessage) error {
	a.sup.setWaiting(a.id, true)
	defer a.sup.setWaiting(a.id, false)
	return a.sup.opts.Policy.ApproveAs(ctx, a.id, tool, input)
}

func (s *Supervisor) setWaiting(id string, waiting bool) {
	s.mu.Lock()
	c := s.byID[id]
	if c == nil {
		s.mu.Unlock()
		return
	}
	switch {
	case waiting && c.info.Status == Running:
		c.info.Status = Waiting
	case !waiting && c.info.Status == Waiting:
		c.info.Status = Running
	default:
		// Paused, killed, or already finished: an approval coming or going
		// does not get to overwrite that.
		s.mu.Unlock()
		return
	}
	s.mu.Unlock()
	s.changed()
}

func addUsage(a, b llm.Usage) llm.Usage {
	a.InputTokens += b.InputTokens
	a.OutputTokens += b.OutputTokens
	a.CacheReadInputTokens += b.CacheReadInputTokens
	a.CacheCreationInputTokens += b.CacheCreationInputTokens
	return a
}