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"`
	Assertions []assertion `yaml:"assertions"`
}

// 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_once    — propose_plan was called exactly once
//	no_duplicate_proposal — propose_plan called at most once
//	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
//	no_rerun              — `run` was NOT called after the followup turn (if any)
//	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_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 "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
}