// Command evals runs the task suite and prints the numbers that go in
// METRICS.md. Everything after Phase 3 gets measured; everything before it was
// guesswork.
//
//	go run ./evals -n 3
package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

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

// Task is one directory under evals/tasks.
type Task struct {
	Name   string
	Dir    string
	Prompt string
}

// Result is one run of one task. Persisted verbatim so trajectories can be read
// back later; most regressions are obvious within ten seconds of reading one.
type Result struct {
	Task     string        `json:"task"`
	Run      int           `json:"run"`
	Model    string        `json:"model"`
	Pass     bool          `json:"pass"`
	Turns    int           `json:"turns"`
	Usage    llm.Usage     `json:"usage"`
	WallMS   int64         `json:"wall_ms"`
	Error    string        `json:"error,omitempty"`
	Answer   string        `json:"answer"`
	Verify   string        `json:"verify_output,omitempty"`
	Messages []llm.Message `json:"messages"`
	Events   []metricEvent `json:"events"`
}

type metricEvent struct {
	Type string `json:"type"`
	Tool string `json:"tool,omitempty"`
	OK   *bool  `json:"ok,omitempty"`
	Err  string `json:"err,omitempty"`
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, "evals:", err)
		os.Exit(1)
	}
}

func run() error {
	n := flag.Int("n", 1, "how many times to run each task; agents are nondeterministic, so one run tells you nothing")
	only := flag.String("tasks", "", "comma-separated task names to run (default: all)")
	parallel := flag.Int("parallel", 4, "how many runs at once")
	model := flag.String("model", "", "model id")
	effort := flag.String("effort", "", "output_config.effort")
	maxTurns := flag.Int("max-turns", 40, "turn cap per run")
	contextWindow := flag.Int("context-window", agent.DefaultContextWindow, "model context window")
	noCache := flag.Bool("no-cache", false, "disable prompt cache breakpoints, to measure what they buy")
	noStream := flag.Bool("no-stream", false, "use the non-streaming endpoint")
	timeout := flag.Duration("timeout", 10*time.Minute, "wall-clock cap per run")
	tasksDir := flag.String("dir", "evals/tasks", "directory of task definitions")
	outDir := flag.String("out", "evals/runs", "where to persist transcripts")
	keep := flag.Bool("keep", false, "keep the temporary work directories")
	flag.Parse()

	tasks, err := loadTasks(*tasksDir, *only)
	if err != nil {
		return err
	}
	if len(tasks) == 0 {
		return fmt.Errorf("no tasks found in %s", *tasksDir)
	}
	client := llm.New(llm.WithModel(*model))
	if !client.HasCredentials() {
		return fmt.Errorf("no credentials: set ANTHROPIC_API_KEY")
	}

	stamp := time.Now().UTC().Format("20060102T150405Z")
	runRoot := filepath.Join(*outDir, stamp)
	if err := os.MkdirAll(runRoot, 0o755); err != nil {
		return err
	}
	fmt.Printf("%d tasks × %d runs · %s · transcripts in %s\n\n", len(tasks), *n, client.Model(), runRoot)

	type job struct {
		task Task
		run  int
	}
	var jobs []job
	for _, t := range tasks {
		for i := 1; i <= *n; i++ {
			jobs = append(jobs, job{task: t, run: i})
		}
	}

	var (
		mu      sync.Mutex
		results []Result
		sem     = make(chan struct{}, *parallel)
		wg      sync.WaitGroup
	)
	for _, j := range jobs {
		wg.Add(1)
		go func(j job) {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()

			res := runOne(j.task, j.run, client, agent.Options{
				MaxTurns:      *maxTurns,
				Effort:        *effort,
				ContextWindow: *contextWindow,
				NoCache:       *noCache,
				Stream:        !*noStream,
			}, *timeout, *keep)
			mu.Lock()
			results = append(results, res)
			mu.Unlock()

			status := "FAIL"
			if res.Pass {
				status = "pass"
			}
			fmt.Printf("  %-26s run %d  %-4s  %2d turns  %5.1fs  %s\n",
				j.task.Name, j.run, status, res.Turns, float64(res.WallMS)/1000, res.Error)

			if err := persist(runRoot, res); err != nil {
				fmt.Fprintln(os.Stderr, "evals: persisting transcript:", err)
			}
		}(j)
	}
	wg.Wait()

	sort.Slice(results, func(i, j int) bool {
		if results[i].Task != results[j].Task {
			return results[i].Task < results[j].Task
		}
		return results[i].Run < results[j].Run
	})
	report(os.Stdout, results)
	return persistSummary(runRoot, results)
}

// runOne sets up a fresh working directory, drives the agent headless, and
// verifies the result.
// The return value is named so the deferred timing assignment lands in the
// value actually returned rather than in a copy already made.
func runOne(task Task, runNo int, client *llm.Client, opts agent.Options, timeout time.Duration, keep bool) (res Result) {
	res = Result{Task: task.Name, Run: runNo, Model: client.Model()}
	start := time.Now()
	defer func() { res.WallMS = time.Since(start).Milliseconds() }()

	work, err := os.MkdirTemp("", "eval-"+task.Name+"-")
	if err != nil {
		res.Error = err.Error()
		return res
	}
	if !keep {
		defer os.RemoveAll(work)
	}

	if out, err := script(context.Background(), filepath.Join(task.Dir, "setup.sh"), work, task.Dir, 2*time.Minute); err != nil {
		res.Error = "setup failed: " + err.Error()
		res.Verify = out
		return res
	}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	var logBuf bytes.Buffer
	log := metrics.NewWriter(&logBuf, fmt.Sprintf("%s-%d", task.Name, runNo))
	state := agent.NewState()
	reg := tools.NewRegistry(log).Register(
		tools.NewRead(work, state),
		tools.NewGlob(work),
		tools.NewGrep(work),
		tools.NewEdit(work, state),
		tools.NewWrite(work, state),
		tools.NewBash(work),
	)
	a := agent.New(client, reg, state, log, opts)
	// The eval runner is headless and operates on a throwaway temp directory,
	// so there is nothing to approve and no one to ask.
	yolo := approval.New(nil)
	yolo.YOLO = true
	a.Approve = yolo

	answer, runErr := a.Run(ctx, task.Prompt)
	res.Answer = answer
	if runErr != nil {
		res.Error = runErr.Error()
	}
	usage, turns := state.Usage()
	res.Usage, res.Turns = usage, turns
	res.Messages = state.History()
	res.Events = parseEvents(logBuf.Bytes())

	// verify.sh may want to know what the agent claimed, which is the only way
	// to check a task whose correct outcome is "I could not do this".
	answerFile := filepath.Join(work, ".agent-answer.txt")
	os.WriteFile(answerFile, []byte(answer), 0o644)

	verifyOut, verifyErr := script(ctx, filepath.Join(task.Dir, "verify.sh"), work, task.Dir, 3*time.Minute)
	res.Verify = verifyOut
	res.Pass = verifyErr == nil
	return res
}

// script runs one of a task's shell files with the work directory as cwd.
func script(ctx context.Context, path, work, taskDir string, timeout time.Duration) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// The script path must be absolute: cmd.Dir is the work directory, so a
	// path relative to the runner's cwd would not resolve.
	cmd := exec.CommandContext(ctx, "bash", mustAbs(path))
	cmd.Dir = work
	cmd.Env = append(os.Environ(),
		"WORK_DIR="+work,
		"TASK_DIR="+mustAbs(taskDir),
		"AGENT_ANSWER_FILE="+filepath.Join(work, ".agent-answer.txt"),
	)
	out, err := cmd.CombinedOutput()
	return string(out), err
}

func mustAbs(p string) string {
	abs, err := filepath.Abs(p)
	if err != nil {
		return p
	}
	return abs
}

func loadTasks(dir, only string) ([]Task, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	wanted := map[string]bool{}
	for _, name := range strings.Split(only, ",") {
		if name = strings.TrimSpace(name); name != "" {
			wanted[name] = true
		}
	}

	var tasks []Task
	for _, e := range entries {
		if !e.IsDir() || (len(wanted) > 0 && !wanted[e.Name()]) {
			continue
		}
		taskDir := filepath.Join(dir, e.Name())
		prompt, err := os.ReadFile(filepath.Join(taskDir, "prompt.txt"))
		if err != nil {
			return nil, fmt.Errorf("%s: %w", e.Name(), err)
		}
		for _, required := range []string{"setup.sh", "verify.sh"} {
			if _, err := os.Stat(filepath.Join(taskDir, required)); err != nil {
				return nil, fmt.Errorf("%s: missing %s", e.Name(), required)
			}
		}
		tasks = append(tasks, Task{Name: e.Name(), Dir: taskDir, Prompt: strings.TrimSpace(string(prompt))})
	}
	return tasks, nil
}

func parseEvents(b []byte) []metricEvent {
	var out []metricEvent
	for _, line := range bytes.Split(b, []byte("\n")) {
		if len(bytes.TrimSpace(line)) == 0 {
			continue
		}
		var ev metricEvent
		if json.Unmarshal(line, &ev) == nil {
			out = append(out, ev)
		}
	}
	return out
}

func persist(runRoot string, res Result) error {
	dir := filepath.Join(runRoot, res.Task)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	b, err := json.MarshalIndent(res, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(dir, fmt.Sprintf("%d.json", res.Run)), b, 0o644)
}

func persistSummary(runRoot string, results []Result) error {
	var buf bytes.Buffer
	report(&buf, results)
	return os.WriteFile(filepath.Join(runRoot, "summary.txt"), buf.Bytes(), 0o644)
}