package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"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) Add(n int) {
	b.mu.Lock()
	b.spent += n
	b.mu.Unlock()
}

type Flock struct {
	brain Brain
	g     Guardrails
	log   *FlightLog
	sem   chan struct{}

	mu   sync.Mutex
	next int
}

func NewFlock(brain Brain, g Guardrails, log *FlightLog) *Flock {
	return &Flock{brain: brain, g: g, log: log, sem: make(chan struct{}, g.MaxCrowsInFlight)}
}

// pick rotates through the roster so every crow gets a chance to build a record.
func (f *Flock) pick(n int) []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]string, n)
	for i := range out {
		out[i] = roster[f.next%len(roster)]
		f.next++
	}
	return out
}

// ask is the only way any crow touches the model: budget, flock cap, timeout, log.
func (f *Flock) ask(ctx context.Context, 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
		}
	}
	f.sem <- struct{}{}
	defer func() { <-f.sem }()

	cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout)
	defer cancel()
	start := time.Now()
	text, tokens, err := f.brain.Ask(cctx, 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, Tokens: tokens, Ms: time.Since(start).Milliseconds(), Output: text}
	if err != nil {
		e.Error = err.Error()
	}
	f.log.Write(e)
	return text, err
}

func (f *Flock) Triage(ctx context.Context, item string, budgets ...*Budget) (Triage, error) {
	text, err := f.ask(ctx, "Sentry", sentrySystem, "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"}
	}
	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, nil
}

func (f *Flock) Execute(ctx context.Context, item string, t Triage, 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(len(t.Subtasks))
		reports, err := f.fanOut(ctx, crows, t.Subtasks, budgets...)
		if err != nil {
			return Outcome{}, err
		}
		return f.reconcile(ctx, item, crows, t.Subtasks, reports, budgets...)
	case "murder":
		n := min(f.g.MurderSize, f.g.MaxCrowsInFlight)
		crows := f.pick(n)
		tasks := make([]string, n)
		for i := range tasks {
			tasks[i] = item
		}
		reports, err := f.fanOut(ctx, crows, tasks, budgets...)
		if err != nil {
			return Outcome{}, err
		}
		return f.reconcile(ctx, item, crows, tasks, reports, budgets...)
	default:
		c := f.pick(1)[0]
		ans, err := f.ask(ctx, c, crowSystem, "TASK: "+item, budgets...)
		return Outcome{Answer: ans, Crows: []string{c}}, err
	}
}

func (f *Flock) fanOut(ctx context.Context, 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, crows[i], crowSystem, "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, 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, "Elder", elderSystem, sb.String(), budgets...)
	return Outcome{Answer: ans, Crows: append(append([]string{}, crows...), "Elder")}, err
}