//! Agent-test data model for a parsed mcptest suite.
//!
//! The `agents:` section types (the agent test itself, its model
//! references, judged eval metrics, and the provider declarations the
//! model matrix resolves against) live here, split out of
//! [`super::types`] so both halves of the suite model stay under the
//! AGENTS module-size ceiling. Everything here is re-exported from
//! [`crate::suite`] so external paths such as
//! `mcptest_core::suite::AgentTest` keep resolving unchanged.

use super::agent_gates::{
    ContextRotProbeGate, DistractorGate, EqualFunctionGate, GoldenPathGate, NameFreeDiscovery,
    OrchestrationGate, RecoveryGate, StabilityGate, TokenEfficiencyGate, ToolEdgesGate,
    ToolSelectionFloor, TrajectoryAxesGate, TrajectoryGate,
};
use super::types::{AssertSet, Assertion, DerivedMetric};

/// User-declared custom provider.
///
/// Today only the OpenAI Chat Completions API shape is supported;
/// every gateway and self-hosted server that targets that shape (Azure
/// OpenAI, OpenRouter, vLLM, llama.cpp, LiteLLM, Together, Groq, Anyscale,
/// Fireworks, etc.) is reachable through one of these.
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderSpec {
    /// Identifier the user writes under `providers:` and references
    /// from `models:` entries via `{provider: <name>, id: ...}`.
    pub name: String,
    /// Wire protocol. The only value today is [`ProviderKind::OpenAi`].
    pub kind: ProviderKind,
    /// Endpoint base URL (`https://my-resource.openai.azure.com/...`,
    /// `https://openrouter.ai/api/v1`, `http://localhost:8000/v1`).
    pub base_url: String,
    /// Environment variable the runtime reads to populate the
    /// `Authorization: Bearer ...` header. Optional so unauthenticated
    /// local endpoints can omit it.
    pub api_key_env: Option<String>,
    /// Optional `OpenAI-Organization` header value, for the few
    /// providers that require it.
    pub organization: Option<String>,
}

/// Wire protocol a [`ProviderSpec`] targets. Extension point for new
/// shapes (Anthropic-on-gateway, raw HuggingFace TGI). Today the only
/// variant is the OpenAI Chat Completions API.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderKind {
    /// Speaks the OpenAI Chat Completions / Responses API.
    OpenAi,
}

/// One agent test.
///
/// An agent test points a real LLM at one or more MCP servers and asserts
/// against the resulting conversation: which tools the model called, what
/// the server returned, the final text reply, and the conversation
/// telemetry (tokens, duration, message count).
#[derive(Debug, Clone, PartialEq)]
pub struct AgentTest {
    /// Human-readable name from the YAML `name:` field.
    pub name: String,
    /// Models to dispatch the same test against. A YAML `model: <id>`
    /// is sugar for `models: [<id>]`; a YAML `models: [<entry>, ...]`
    /// runs the test once per entry so the reporter can show pass /
    /// fail per model. Each entry is either a bare string (auto-detect
    /// the provider family from the prefix) or an object referencing
    /// a custom `providers:` entry. Always non-empty after parsing.
    pub models: Vec<ModelRef>,
    /// Server keys this test makes available to the model. Each entry is
    /// resolved against [`crate::suite::ResolvedSuite::servers`] at run time. The driver
    /// dispatches to each in turn.
    pub servers: Vec<String>,
    /// User prompt that kicks off the conversation.
    pub prompt: String,
    /// Optional system prompt. Defaults to none.
    pub system_prompt: Option<String>,
    /// Hard cap on the number of tool-use iterations. Defaults to 10
    /// when the YAML omits the field.
    pub max_turns: u32,
    /// Per-call `max_tokens` for the model. Defaults to 1024.
    pub max_tokens: u32,
    /// Cumulative token budget across the whole run. Distinct
    /// from per-call [`Self::max_tokens`]. `None` (the YAML default) leaves
    /// the run unbounded so existing suites are unchanged.
    pub token_budget: Option<u32>,
    /// Cap on total tool invocations across the run. `None` (the
    /// YAML default) leaves the run unbounded.
    pub max_tool_calls: Option<u32>,
    /// Wall-clock deadline for the whole run as a [`std::time::Duration`]
    /// parsed from a human string like `"30s"` or `"2m"`. `None`
    /// (the YAML default) leaves the run unbounded.
    pub time_budget: Option<std::time::Duration>,
    /// Assertions evaluated against the [`crate::suite::AgentTest`]'s
    /// trace envelope. See the agent target grammar in the crate docs.
    pub expectations: Vec<Assertion>,
    /// Judged `eval:` metrics scored over the run trace after it completes
    ///. Each produces a pass/fail report row. Empty by default.
    pub eval: Vec<EvalMetric>,
    /// Optional multi-turn conversation surface. `None` (the YAML
    /// default) drives the single `prompt`; `Some` replays a scripted turn
    /// list or runs an LLM-driven simulated user. The per-turn trace lands in
    /// `conversation.per_turn` and feeds `multi_turn_mcp_use`.
    pub turns: Option<AgentTurns>,
    /// Optional tags from the YAML `tags: [...]` field. The
    /// same tag applies to every (agent, model) row the matrix
    /// expands to, so `--tag smoke` keeps or drops every model for an
    /// agent test as a unit.
    pub tags: Vec<String>,
    /// Optional `skip: "<reason>"` marker (WOR-1267). Applies to every
    /// (agent, model) cell the matrix expands to: each is reported as
    /// SKIP with this reason and no model is dispatched.
    pub skip: Option<String>,
    /// `only: true` focus marker (WOR-1267). Marks every cell of this
    /// agent test as focused; see [`crate::suite::TypedTest::only`] for the run
    /// semantics shared by all test kinds.
    pub only: bool,
    /// How many times each (agent, model) pair runs. Defaults
    /// to 1 (single run, legacy behavior). Values above 1 drive the
    /// multi-run pass^k aggregation at the CLI level. Always >= 1 after
    /// parsing; a YAML `runs: 0` is clamped up to 1.
    pub runs: u32,
    /// Optional cross-run tool-selection floor. When set, the
    /// CLI aggregates the N run traces into a [`crate::eval::selection::SelectionReport`]
    /// and fails the run if the floor is not met. `None` means the
    /// agent test only asserts via per-run `expect:` blocks.
    pub tool_selection: Option<ToolSelectionFloor>,
    /// Optional comparative matrix assertion. Unlike a per-row
    /// `threshold:`, this spans every cell of the model matrix this agent
    /// test fans out to: it identifies the best-scoring cell across all
    /// models and can gate the run on a `max_score` floor. `None` means the
    /// agent test makes no cross-cell comparison.
    pub matrix_assert: Option<MatrixAssert>,
    /// Optional within-session stability gate. When set, the CLI
    /// aggregates the N run traces into a
    /// [`crate::eval::stability_aggregate::StabilitySummary`] and gates its
    /// `stability.score`, `stability.weakest_score`, and `stability.variance`
    /// targets. `None` means the agent test makes no stability assertion.
    pub stability: Option<StabilityGate>,
    /// Optional objective tool-selection F1 gate over declared equal-function
    /// sets (MSC-Bench). When set, the CLI folds the N run traces into the
    /// `tool_selection.f1`, `tool_selection.precision`, and
    /// `tool_selection.recall` targets and gates on them with no model in the
    /// loop. `None` means the agent test makes no F1 assertion.
    pub equal_function: Option<EqualFunctionGate>,
    /// Optional tool-surface token-efficiency gate (WOR-1366). When set, the CLI
    /// scores the same equal-function selection F1 and folds in the serialized
    /// tool-surface token count and the run cost into the `token_efficiency.*`
    /// targets (`f1`, `tokens_per_correct`, `cost_per_correct`, `grade`, ...),
    /// gating on them with no extra model call. `None` means the agent test
    /// makes no efficiency assertion. Live-run only: a cassette replay carries no
    /// catalog token count, so the gate does not fire on replay.
    pub token_efficiency: Option<TokenEfficiencyGate>,
    /// Optional structural tool-edge coverage gate (WOR-1242). When set, the
    /// CLI folds the run trace against the declared allowed, restricted, and
    /// delegation edges into the `edges.*` targets and gates on them, with no
    /// model in the loop. `None` means the agent test makes no edge assertion.
    pub tool_edges: Option<ToolEdgesGate>,
    /// Optional orchestration diagnostics gate (MCP-Atlas). When set, the CLI
    /// folds the run trace into the five `orchestration.*` sub-score targets
    /// (discovery, parameterization, syntax, error_recovery, efficiency) and
    /// gates on them with no model in the loop. `None` means the agent test
    /// makes no orchestration assertion.
    pub orchestration: Option<OrchestrationGate>,
    /// Optional name-free discovery declaration (MCP-Atlas). When `Some` with
    /// `name_free` true, the prompt is intent-only and names no tools or
    /// servers, so the discovered path is scored against equal-function
    /// classes. `None` means the prompt may name tools as usual.
    pub discovery: Option<NameFreeDiscovery>,
    /// Optional distractor-tool / tool-overload gate (MCPAgentBench). `None`
    /// when the agent test declares no `distractors:` block. The runner injects
    /// the declared distractors into the candidate list (preview) and scores the
    /// recorded selection into the `distractors.accuracy` and
    /// `distractors.chose_distractor` targets with no model in the loop.
    pub distractors: Option<DistractorGate>,
    /// Optional context-rot knee probe gate (context-rot sweep). `None` when the
    /// agent test declares no `context_rot_probe:` block. When set, the runner
    /// sweeps the served-catalog size across the declared levels (injecting that
    /// many distractor tools per level), re-runs the agent at each level, scores
    /// per-level selection accuracy, and folds the `(context_tokens, quality)`
    /// points into the `context_rot.*` targets. Live-run only: the sweep re-runs
    /// the agent at each level, so a cassette replay does not fire it.
    pub context_rot_probe: Option<ContextRotProbeGate>,
    /// Optional golden-path efficiency gate. `None` when the agent test
    /// declares no `golden_path:` block. The runner scores the recorded call
    /// sequence against the declared ideal path into the `golden_path.*`
    /// targets with no model in the loop.
    pub golden_path: Option<GoldenPathGate>,
    /// Optional trajectory validation gate. `None` when the agent test declares
    /// no `trajectory:` block. The runner validates the recorded call sequence
    /// against the expected sequence under the chosen match mode into the
    /// `trajectory.*` targets with no model in the loop.
    pub trajectory: Option<TrajectoryGate>,
    /// Optional four-axis trajectory scorecard gate (WOR-1368). `None` when the
    /// agent test declares no `trajectory_axes:` block. The runner scores the
    /// recorded call sequence against the declared dependency and order
    /// constraints into the `trajectory.dependency_satisfaction` and
    /// `trajectory.order_satisfaction` targets with no model in the loop.
    pub trajectory_axes: Option<TrajectoryAxesGate>,
    /// Names of suite-level `faults:` this agent run injects. Empty (the YAML
    /// default) means no fault is injected. Each name is validated against the
    /// suite's `faults:` block at parse time.
    pub inject: Vec<String>,
    /// Optional `recovery:` gate. When `Some`, the executor injects the named
    /// faults, synthesizes the unresponsive call in virtual time, and gates the
    /// run on the recovery budget. `None` means the agent test makes no recovery
    /// assertion.
    pub recovery: Option<RecoveryGate>,
    /// Named assert-set bundles. Empty when none are declared.
    pub assert_sets: Vec<AssertSet>,
    /// Optional combined-score threshold on a 0..1 scale. When
    /// `None`, every assertion and assert-set must pass (legacy
    /// behavior); when set, the test passes on a combined weighted score
    /// at or above the threshold.
    pub threshold: Option<f64>,
    /// Derived metrics reported for this agent test. Empty
    /// when none are declared.
    pub derived_metrics: Vec<DerivedMetric>,
}

/// A judged `eval:` metric on an agent test. Each variant scores the
/// run trace 0..1 through the provider and gates on a threshold. The enum grows
/// as the sibling metrics land (task-completion, mcp-use, argument-correctness).
#[derive(Debug, Clone, PartialEq)]
pub enum EvalMetric {
    /// A weighted, multi-criterion rubric.
    Rubric(crate::eval::rubric::RubricConfig),
    /// Did the agent accomplish the task (the prompt)? A single judged
    /// criterion over the run trace.
    TaskCompletion {
        /// Pass threshold on the 0..1 score.
        threshold: f64,
        /// When true, pass only on a perfect score.
        strict: bool,
    },
    /// Did the agent use the right tools with the right arguments for the
    /// intent? A single judged criterion over the run trace.
    McpUse {
        /// Pass threshold on the 0..1 score.
        threshold: f64,
        /// When true, pass only on a perfect score.
        strict: bool,
    },
    /// Were the tool-call arguments correct for the intent? A single judged
    /// criterion focused on argument fidelity.
    ArgumentCorrectness {
        /// Pass threshold on the 0..1 score.
        threshold: f64,
        /// When true, pass only on a perfect score.
        strict: bool,
    },
    /// Was the sequence of tool calls a sensible plan (no redundant or
    /// misordered calls, no missing prerequisite)? A single judged criterion.
    PlanQuality {
        /// Pass threshold on the 0..1 score.
        threshold: f64,
        /// When true, pass only on a perfect score.
        strict: bool,
    },
    /// MCP-use alignment averaged across a multi-turn conversation.
    /// Each `conversation.per_turn` slice is judged for tool-use quality and
    /// the per-turn scores are averaged. Requires a `turns:` surface; with no
    /// per-turn trace the metric defers.
    MultiTurnMcpUse {
        /// Pass threshold on the 0..1 aggregate.
        threshold: f64,
        /// When true, pass only on a perfect score.
        strict: bool,
    },
}

/// How a multi-turn agent conversation is driven.
///
/// `None` on [`AgentTest::turns`] is the legacy single-`prompt` run. A
/// `Scripted` list replays fixed user messages in order; `Simulated` lets an
/// LLM play the user from a goal and persona.
#[derive(Debug, Clone, PartialEq)]
pub enum AgentTurns {
    /// An ordered list of scripted user messages the driver replays.
    Scripted(Vec<String>),
    /// An LLM-driven simulated user.
    Simulated(SimulatedUser),
}

/// Config for an LLM-driven simulated user.
#[derive(Debug, Clone, PartialEq)]
pub struct SimulatedUser {
    /// What the simulated user is trying to accomplish.
    pub goal: String,
    /// Optional persona (tone, behavior) the simulator adopts.
    pub persona: Option<String>,
    /// Cap on the number of user turns the simulator generates. Defaults to 6
    /// when the YAML omits it, so a runaway simulation stays bounded.
    pub max_turns: u32,
}

/// Comparative assertion across a model-matrix fan-out.
///
/// A single agent test can run against several models (`models: [...]`),
/// producing one "cell" per model. The per-cell `threshold:` judges each
/// cell on its own; a `matrix_assert` instead compares cells against each
/// other. Two independent constructs:
///
/// - `max_score`: the gating one. The run fails unless the BEST cell's
///   combined score across the whole matrix is at or above this floor. Use
///   it to assert "at least one model can do this well enough."
/// - `select_best`: reporting only. Identifies and prints the
///   highest-scoring cell across the matrix; it never fails the run.
///
/// Both depend on a per-cell combined score, so declaring either makes
/// every cell of this agent test carry a `report.tests[].score` even when
/// the cell sets no `threshold:`. A test that declares neither key keeps
/// the legacy report JSON unchanged.
#[derive(Debug, Clone, PartialEq)]
pub struct MatrixAssert {
    /// Minimum combined score the best matrix cell must reach for the run
    /// to pass. `None` means no gate (use `select_best` for reporting only).
    pub max_score: Option<f64>,
    /// When `true`, print the highest-scoring cell across the matrix. Pure
    /// reporting: it never changes the run's pass/fail.
    pub select_best: bool,
}

impl MatrixAssert {
    /// True when the agent test wants its cells to carry a score even
    /// without a per-cell `threshold:`. Either construct (the gate or the
    /// report line) needs a per-cell score to compare, so any declared
    /// `matrix_assert` opts its cells into score emission.
    pub fn wants_score(&self) -> bool {
        self.max_score.is_some() || self.select_best
    }
}

/// How an agent test names a model. Bare strings are auto-detected
/// by the provider family prefix (`claude-*` -> Anthropic env var,
/// `gpt-*` -> OpenAI env var, etc); object entries reference a
/// custom `providers:` block by name. .
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelRef {
    /// Bare model id; provider family inferred from the prefix.
    AutoDetect(String),
    /// Explicit named-provider reference. The `provider` must match a
    /// key in [`crate::suite::ResolvedSuite::providers`]; `id` is the model
    /// identifier passed to the gateway.
    Named {
        /// Provider key declared under `providers:`.
        provider: String,
        /// Model id the gateway expects (`anthropic/claude-3.5-sonnet`,
        /// `openai/gpt-4o`, `meta-llama/Llama-3.1-70B-Instruct`).
        id: String,
    },
}

impl ModelRef {
    /// Render the entry as a human-readable label used in reporter
    /// case names and tracing logs. Auto-detected entries show only
    /// the id; named entries show `<provider>/<id>` so the reader can
    /// tell which gateway was hit.
    pub fn display_label(&self) -> String {
        match self {
            ModelRef::AutoDetect(id) => id.clone(),
            ModelRef::Named { provider, id } => format!("{provider}/{id}"),
        }
    }

    /// Model id the provider sees on the wire. For named refs this is
    /// the `id:` field verbatim; for auto-detect it is the bare string.
    pub fn model_id(&self) -> &str {
        match self {
            ModelRef::AutoDetect(id) => id,
            ModelRef::Named { id, .. } => id,
        }
    }
}

/// Top-level budget knob shared by every agent test in a suite.
///
/// Two values, both in cents. `None` for either field means "no cap" for
/// that scope. The runner threads one [`SuiteBudget`] through every
/// agent test in the file.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct SuiteBudget {
    /// Per-test spend cap in USD cents. The runner resets the per-test
    /// counter between tests.
    pub per_test_usd_cents: Option<u64>,
    /// Per-suite spend cap in USD cents. The runner aggregates spend
    /// across every agent test in the file.
    pub per_suite_usd_cents: Option<u64>,
}
