package main

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

const (
	markerSentry = "You are the Sentry"
	markerElder  = "You are the Elder"

	sentrySystem = markerSentry + ` of a flock of AI agents called crows.
Classify one todo item. Reply with JSON only:
{"mode":"solo|split|murder","task_type":"retrieval|writing|analysis|other","subtasks":["..."],"reason":"one short sentence"}
- solo: simple; one crow can answer it.
- split: large; break into 2-4 independent subtasks (fill "subtasks").
- murder: ambiguous or high-stakes; several crows answer independently, then reconcile.
Crows are text-only: no web, no files, no tools.`

	crowSystem = `You are a crow in a flock of AI agents. Complete the task concisely.
You are text-only: no web, no files, no tools. If the task needs something you cannot do, say so plainly instead of guessing.`

	elderSystem = markerElder + ` crow. Several crows reported on one task.
Merge their reports into one answer. Where they disagree, say so and say which view is better supported. Be concise.`
)

var roster = []string{"Hugin", "Munin", "Corvina", "Soot", "Pepper", "Inky", "Edgar", "Jet"}

type Triage struct {
	Mode     string   `json:"mode"`
	TaskType string   `json:"task_type"`
	Subtasks []string `json:"subtasks,omitempty"`
	Reason   string   `json:"reason"`
}

type Outcome struct {
	Answer string
	Crows  []string // everyone accountable for the result
}

type Budget struct {
	mu           sync.Mutex
	name         string
	limit, spent int
}

func (b *Budget) Check() error {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.spent >= b.limit {
		return fmt.Errorf("%s budget spent (%d/%d tokens)", b.name, b.spent, b.limit)
	}
	return nil
}

func (b *Budget) Spent() int {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.spent
}

func (b *Budget) Add(n int) {
	b.mu.Lock()
	b.spent += n
	b.mu.Unlock()
}

// Watcher hears every takeoff and landing (Sentry and Elder included). It may be nil.
type Watcher interface {
	Flying(crow string)
	Landed(crow string, tokens int, err error)
}

type Flock struct {
	brain Brain
	jev   *JevClient // optional classifier for triage; nil means the LLM Sentry triages alone
	g     Guardrails
	log   *FlightLog
	sem   chan struct{}

	mu     sync.Mutex
	next   int            // rotation through the whole flock (the todo nest)
	core   map[string]int // per nest: rotation within its core
	borrow map[string]int // per nest: rotation through crows outside its core
}

func NewFlock(brain Brain, g Guardrails, log *FlightLog) *Flock {
	return &Flock{
		brain: brain, g: g, log: log, sem: make(chan struct{}, g.MaxCrowsInFlight),
		core: map[string]int{}, borrow: map[string]int{},
	}
}

// pick chooses n crows for a job in a nest. A nest with no core (todo) rotates
// through the whole roster so every crow gets a chance to build a record.
// Otherwise core crows go first (rotating within the core), and any extra
// crows are borrowed from outside the core (rotating).
func (f *Flock) pick(nest Nest, n int) []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]string, 0, n)
	if len(nest.Crows) == 0 {
		for range n {
			out = append(out, roster[f.next%len(roster)])
			f.next++
		}
		return out
	}
	core := nest.Crows
	for range min(n, len(core)) {
		out = append(out, core[f.core[nest.ID]%len(core)])
		f.core[nest.ID]++
	}
	outside := slices.DeleteFunc(slices.Clone(roster), func(c string) bool { return slices.Contains(core, c) })
	if len(outside) == 0 {
		outside = core // the whole flock is the core: nobody left to borrow
	}
	for len(out) < n {
		out = append(out, outside[f.borrow[nest.ID]%len(outside)])
		f.borrow[nest.ID]++
	}
	return out
}

// reserve takes a flock slot, honoring ctx: a saturated flock waits for a
// slot, but gives up when the context ends instead of wedging the caller.
func (f *Flock) reserve(ctx context.Context) (func(), bool) {
	select {
	case f.sem <- struct{}{}:
		return func() { <-f.sem }, true
	case <-ctx.Done():
		return nil, false
	}
}

// ask is the only way any crow touches the model: budget, flock cap, timeout, log.
func (f *Flock) ask(ctx context.Context, w Watcher, crow, system, prompt string, budgets ...*Budget) (string, error) {
	for _, b := range budgets {
		if err := b.Check(); err != nil {
			f.log.Write(Event{Kind: "grounded", Crow: crow, Error: err.Error()})
			return "", err
		}
	}
	release, slot := f.reserve(ctx)
	if !slot {
		return "", fmt.Errorf("flock at cap: %w", ctx.Err())
	}
	defer release()

	if w != nil {
		w.Flying(crow)
	}
	cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout)
	defer cancel()
	start := time.Now()
	var model string
	text, tokens, err := f.brain.Ask(withModelSlot(cctx, &model), system, prompt)
	for _, b := range budgets {
		b.Add(tokens)
	}
	if errors.Is(err, context.DeadlineExceeded) {
		err = fmt.Errorf("%s timed out after %s and was dropped", crow, f.g.CrowTimeout)
	}
	e := Event{Kind: "flight", Crow: crow, Model: model, Tokens: tokens, Ms: time.Since(start).Milliseconds(), Output: text}
	if err != nil {
		e.Error = err.Error()
	}
	f.log.Write(e)
	if w != nil {
		w.Landed(crow, tokens, err)
	}
	return text, err
}

// SetClassifier pairs the Sentry with Jev. Pass nil to triage with the LLM alone.
func (f *Flock) SetClassifier(j *JevClient) { f.jev = j }

// Triage: Jev classifies first when it's configured. Its call stands when it is
// confident and no subtasks are needed; otherwise the LLM Sentry decides, and
// the reason records what Jev thought. Jev failing never blocks a list.
func (f *Flock) Triage(ctx context.Context, w Watcher, item string, budgets ...*Budget) (Triage, error) {
	if f.jev == nil {
		return f.llmTriage(ctx, w, item, budgets...)
	}
	call, err := f.classify(ctx, w, item, budgets...)
	if err != nil {
		t, lerr := f.llmTriage(ctx, w, item, budgets...)
		t.Reason = "Jev was unavailable; the LLM Sentry decided. " + t.Reason
		return t, lerr
	}
	if call.confident(f.jev.threshold) && call.Mode != "split" {
		return f.normalize(Triage{Mode: call.Mode, TaskType: call.TaskType, Reason: call.summary()}), nil
	}
	t, lerr := f.llmTriage(ctx, w, item, budgets...)
	switch {
	case !call.confident(f.jev.threshold):
		t.Reason = fmt.Sprintf("%s; unsure, so the LLM Sentry decided. %s", call.summary(), t.Reason)
	case t.Mode == "split":
		t.Reason = fmt.Sprintf("%s; the LLM Sentry wrote the subtasks. %s", call.summary(), t.Reason)
	default:
		t.Reason = fmt.Sprintf("%s; the LLM Sentry disagreed and chose %s. %s", call.summary(), t.Mode, t.Reason)
	}
	return t, lerr
}

func (f *Flock) llmTriage(ctx context.Context, w Watcher, item string, budgets ...*Budget) (Triage, error) {
	text, err := f.ask(ctx, w, "Sentry", systemFor("Sentry", sentrySystem, false), "ITEM: "+item, budgets...)
	if err != nil {
		return Triage{}, err
	}
	var t Triage
	if i, j := strings.Index(text, "{"), strings.LastIndex(text, "}"); i < 0 || j < i || json.Unmarshal([]byte(text[i:j+1]), &t) != nil {
		t = Triage{Mode: "solo", TaskType: "other", Reason: "Sentry's call was unreadable; defaulting to solo"}
	}
	return f.normalize(t), nil
}

// normalize keeps any triage, from Jev or the LLM, inside the guardrails.
func (f *Flock) normalize(t Triage) Triage {
	t.Mode = strings.ToLower(strings.TrimSpace(t.Mode))
	t.TaskType = strings.ToLower(strings.TrimSpace(t.TaskType))
	switch t.TaskType {
	case "retrieval", "writing", "analysis", "other":
	default:
		t.TaskType = "other"
	}
	switch t.Mode {
	case "split":
		if len(t.Subtasks) > f.g.MaxSubtasks {
			t.Subtasks = t.Subtasks[:f.g.MaxSubtasks]
		}
		if len(t.Subtasks) < 2 {
			t.Mode, t.Subtasks = "solo", nil
		}
	case "murder":
		t.Subtasks = nil
	default:
		t.Mode, t.Subtasks = "solo", nil
	}
	return t
}

func (f *Flock) Execute(ctx context.Context, w Watcher, item string, t Triage, nest Nest, budgets ...*Budget) (Outcome, error) {
	switch t.Mode {
	case "split":
		// Split depth is 1: subtasks go straight to crows, never back to the Sentry.
		crows := f.pick(nest, len(t.Subtasks))
		reports, err := f.fanOut(ctx, w, crows, t.Subtasks, budgets...)
		if err != nil {
			return Outcome{}, err
		}
		return f.reconcile(ctx, w, item, crows, t.Subtasks, reports, budgets...)
	case "murder":
		n := min(f.g.MurderSize, f.g.MaxCrowsInFlight)
		crows := f.pick(nest, n)
		tasks := make([]string, n)
		for i := range tasks {
			tasks[i] = item
		}
		reports, err := f.fanOut(ctx, w, crows, tasks, budgets...)
		if err != nil {
			return Outcome{}, err
		}
		return f.reconcile(ctx, w, item, crows, tasks, reports, budgets...)
	default:
		c := f.pick(nest, 1)[0]
		ans, err := f.ask(ctx, w, c, systemFor(c, crowSystem, true), "TASK: "+item, budgets...)
		return Outcome{Answer: ans, Crows: []string{c}}, err
	}
}

func (f *Flock) fanOut(ctx context.Context, w Watcher, crows, tasks []string, budgets ...*Budget) ([]string, error) {
	reports := make([]string, len(tasks))
	errs := make([]error, len(tasks))
	var wg sync.WaitGroup
	for i := range tasks {
		wg.Add(1)
		go func() {
			defer wg.Done()
			reports[i], errs[i] = f.ask(ctx, w, crows[i], systemFor(crows[i], crowSystem, true), "TASK: "+tasks[i], budgets...)
		}()
	}
	wg.Wait()
	ok := 0
	for i, err := range errs {
		if err != nil {
			reports[i] = "(no report: " + err.Error() + ")"
		} else {
			ok++
		}
	}
	if ok == 0 {
		return nil, fmt.Errorf("no crow came back: %w", errs[0])
	}
	return reports, nil
}

func (f *Flock) reconcile(ctx context.Context, w Watcher, item string, crows, tasks, reports []string, budgets ...*Budget) (Outcome, error) {
	var sb strings.Builder
	fmt.Fprintf(&sb, "ORIGINAL TASK: %s\n\n", item)
	for i := range reports {
		fmt.Fprintf(&sb, "--- %s on %q:\n%s\n\n", crows[i], tasks[i], reports[i])
	}
	ans, err := f.ask(ctx, w, "Elder", systemFor("Elder", elderSystem, true), sb.String(), budgets...)
	return Outcome{Answer: ans, Crows: append(append([]string{}, crows...), "Elder")}, err
}