package main

import (
	"fmt"
	"os"

	"gopkg.in/yaml.v3"
)

// conversation is one golden conversation from a manifest.
type conversation struct {
	Name       string     `yaml:"name"`
	Prompt     string     `yaml:"prompt"`
	Followup   string     `yaml:"followup"`   // backward compat: single followup
	Followups  []string   `yaml:"followups"`  // P5: multi-turn followups
	Assertions []assertion `yaml:"assertions"`
}

// followups returns the full list of follow-up messages, supporting both
// the single `followup` field (backward compat) and the multi-turn
// `followups` list.
func (c conversation) followups() []string {
	if len(c.Followups) > 0 {
		return c.Followups
	}
	if c.Followup != "" {
		return []string{c.Followup}
	}
	return nil
}

// assertion is one check against the final transcript. The `kind` field
// selects the scorer; the rest are scorer-specific parameters.
//
// Supported kinds:
//
//	completes             — session status reached done/failed (not stuck executing)
//	outcome_is            — session outcome == value (success/failure/partial)
//	no_propose_plan       — propose_plan was never called
//	proposes_plan         — propose_plan called >= 1 time (plan-always model; P1)
//	proposes_plan_once    — propose_plan was called exactly once
//	no_duplicate_proposal — propose_plan called at most once
//	plan_before_run       — the first `run` call comes after the first `propose_plan` (P1 ordering gate)
//	plan_generations      — the persisted plan has exactly `value` distinct generations (P2 iteration: 1 = single, 2 = one followup)
//	writes_back           — update_entity_attributes or create_relationship was called
//	max_tool_calls        — total tool calls <= value
//	max_run_calls         — total `run` calls <= value
//	no_run                — `run` was never called
//	calls_tool            — the named tool appears in the transcript
//	plan_step_count       — the plan has exactly `value` steps
//	no_duplicate_complete — complete_task called at most once
type assertion struct {
	Kind  string `yaml:"kind"`
	Value any    `yaml:"value"`
}

// loadManifest reads a YAML file containing a list of conversations.
func loadManifest(path string) ([]conversation, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var convs []conversation
	if err := yaml.Unmarshal(b, &convs); err != nil {
		return nil, fmt.Errorf("parse %s: %w", path, err)
	}
	return convs, nil
}

// scoreAssertions evaluates each assertion against the transcript + session.
func scoreAssertions(asserts []assertion, t transcript, s sessionState) []assertionResult {
	out := make([]assertionResult, 0, len(asserts))
	for _, a := range asserts {
		r := assertionResult{Name: a.Kind}
		r.Passed, r.Detail = scoreOne(a, t, s)
		if !r.Passed && r.Detail == "" {
			r.Detail = "assertion failed"
		}
		out = append(out, r)
	}
	return out
}

func scoreOne(a assertion, t transcript, s sessionState) (bool, string) {
	tools := t.toolNames()
	switch a.Kind {
	case "completes":
		if s.Status == "done" || s.Status == "failed" {
			return true, fmt.Sprintf("status=%s", s.Status)
		}
		return false, fmt.Sprintf("status=%s (not terminal)", s.Status)

	case "outcome_is":
		want, _ := a.Value.(string)
		if s.Outcome == want {
			return true, fmt.Sprintf("outcome=%s", s.Outcome)
		}
		return false, fmt.Sprintf("outcome=%s, want %s", s.Outcome, want)

	case "no_propose_plan":
		n := countTool(tools, "propose_plan")
		if n == 0 {
			return true, "propose_plan not called"
		}
		return false, fmt.Sprintf("propose_plan called %d time(s)", n)

	case "proposes_plan":
		// P1 plan-always: propose_plan called >= 1 time.
		n := countTool(tools, "propose_plan")
		if n >= 1 {
			return true, fmt.Sprintf("propose_plan called %d time(s)", n)
		}
		return false, "propose_plan never called (plan-always requires >= 1)"

	case "proposes_plan_once":
		n := countTool(tools, "propose_plan")
		if n == 1 {
			return true, "propose_plan called once"
		}
		return false, fmt.Sprintf("propose_plan called %d time(s), want 1", n)

	case "no_duplicate_proposal":
		n := countTool(tools, "propose_plan")
		if n <= 1 {
			return true, fmt.Sprintf("propose_plan called %d time(s)", n)
		}
		return false, fmt.Sprintf("propose_plan called %d time(s), want <= 1", n)

	case "plan_before_run":
		// P1 ordering gate: the first `run` call's global index in the
		// transcript is strictly greater than the first `propose_plan`
		// index. Both indices are over the flat tool-call list (across all
		// messages, in order).
		planIdx, runIdx := -1, -1
		for i, name := range tools {
			if name == "propose_plan" && planIdx == -1 {
				planIdx = i
			}
			if name == "run" && runIdx == -1 {
				runIdx = i
			}
		}
		if runIdx == -1 {
			return true, "run never called (ordering trivially satisfied)"
		}
		if planIdx == -1 {
			return false, "run called but propose_plan never called"
		}
		if planIdx < runIdx {
			return true, fmt.Sprintf("propose_plan at index %d before run at index %d", planIdx, runIdx)
		}
		return false, fmt.Sprintf("run at index %d before propose_plan at index %d", runIdx, planIdx)

	case "plan_generations":
		// P2 iteration: counts distinct `generation` values in
		// session_plan_steps. 1 = single sub-task, 2 = one follow-up
		// sub-task, etc. Requires the plan endpoint to return generation
		// values; the eval fetches /sessions/{id}/plan and passes it via
		// the transcript's PlanSteps field.
		want := toInt(a.Value)
		gens := t.distinctGenerations()
		if gens == want {
			return true, fmt.Sprintf("%d plan generation(s)", gens)
		}
		return false, fmt.Sprintf("%d plan generation(s), want %d", gens, want)

	case "writes_back":
		n := countTool(tools, "update_entity_attributes") + countTool(tools, "create_relationship")
		if n > 0 {
			return true, fmt.Sprintf("%d writeback call(s)", n)
		}
		return false, "no update_entity_attributes or create_relationship calls"

	case "max_tool_calls":
		max := toInt(a.Value)
		if t.toolCallCount() <= max {
			return true, fmt.Sprintf("%d tool calls (<= %d)", t.toolCallCount(), max)
		}
		return false, fmt.Sprintf("%d tool calls, want <= %d", t.toolCallCount(), max)

	case "max_run_calls":
		max := toInt(a.Value)
		n := countTool(tools, "run")
		if n <= max {
			return true, fmt.Sprintf("%d run calls (<= %d)", n, max)
		}
		return false, fmt.Sprintf("%d run calls, want <= %d", n, max)

	case "no_run":
		n := countTool(tools, "run")
		if n == 0 {
			return true, "run not called"
		}
		return false, fmt.Sprintf("run called %d time(s)", n)

	case "calls_tool":
		want, _ := a.Value.(string)
		n := countTool(tools, want)
		if n > 0 {
			return true, fmt.Sprintf("%s called %d time(s)", want, n)
		}
		return false, fmt.Sprintf("%s not called", want)

	case "no_duplicate_complete":
		n := countTool(tools, "complete_task")
		if n <= 1 {
			return true, fmt.Sprintf("complete_task called %d time(s)", n)
		}
		return false, fmt.Sprintf("complete_task called %d time(s), want <= 1", n)

	default:
		return false, fmt.Sprintf("unknown assertion kind: %s", a.Kind)
	}
}

func countTool(names []string, name string) int {
	n := 0
	for _, x := range names {
		if x == name {
			n++
		}
	}
	return n
}

func toInt(v any) int {
	switch x := v.(type) {
	case int:
		return x
	case int64:
		return int(x)
	case float64:
		return int(x)
	}
	return 0
}