//! Aggregate types for the RunLog projection.

use chrono::{DateTime, Utc};
use klieo_core::ids::RunId;
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Run lifecycle status.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum RunStatus {
    /// Run is still in progress.
    Running,
    /// Run finished successfully.
    Completed,
    /// Run failed.
    Failed,
    /// Run was cancelled by the operator or scheduler.
    Cancelled,
}

/// Discriminator for the kind of side effect a `Step` represents.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum StepKind {
    /// LLM completion or chat call.
    LlmCall,
    /// Tool invocation through `ToolInvoker`.
    ToolCall,
    /// Summarizer checkpoint — distinct from a top-level `LlmCall` so
    /// summarizer overhead does not pollute substantive-reasoning cost
    /// and latency attribution.
    SummaryCheckpoint,
    /// Operational-layer event emitted by klieo-ops (supervisor state
    /// changes, gate decisions, governor denials, etc.). The `input`
    /// field of the corresponding [`Step`] carries the raw
    /// `serde_json::Value` payload from `Episode::Ops`.
    OpsEvent,
}

/// One unit of agent work — either an LLM call or a tool call.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Step {
    /// Zero-based index of the step within the run.
    pub idx: u32,
    /// Discriminator.
    pub kind: StepKind,
    /// Tool name (for `ToolCall`) or model name (for `LlmCall`); `None` if unknown.
    pub name: Option<String>,
    /// JSON inputs (tool args, prompt structure).
    pub input: serde_json::Value,
    /// JSON outputs (tool result, completion text).
    pub output: serde_json::Value,
    /// Error message if this step failed (`None` on success).
    pub error: Option<String>,
    /// Wall-clock latency serialised as milliseconds.
    #[serde(with = "duration_millis")]
    #[cfg_attr(feature = "schemars", schemars(with = "u64"))]
    pub latency: Duration,
    /// Tracing span id (stringly-typed for now; richer linkage is a Wave 7 carryover).
    pub span_id: Option<String>,
}

/// Sidecar I/O record for a single `Episode::LlmCall`.
///
/// `klieo_core::Episode::LlmCall` carries only `tokens` + `latency_ms` (the
/// trait surface is frozen at 0.1.0 — see `docs/SEMVER.md`). To make
/// projected `RunLog`s replayable, callers that have access to the original
/// prompt + completion can pass a parallel `&[LlmIo]` slice to
/// [`crate::projector::project_with_llm_io`]: the i-th `LlmIo` is paired with
/// the i-th `Episode::LlmCall` in the episode stream.
///
/// Without this sidecar, projected `Step.input` / `Step.output` for LLM
/// steps remain `Value::Null`, and [`crate::replay::replay`] cannot reproduce
/// the recorded prompt.
///
/// ## Cost-accounting fields
///
/// `provider`, `prompt_tokens`, and `completion_tokens` are optional sidecar
/// fields used by [`crate::projector::project_with_price_table`] to compute a
/// per-run [`Cost`]. `Episode::LlmCall::tokens` is a single `u32` total with
/// no provider info, so callers populate these fields here when they want
/// USD totals on the projected `RunLog`. Each new field carries
/// `#[serde(default)]` so older serialised records still deserialise.
///
/// ## Trust model
///
/// `provider`, `prompt_tokens`, and `completion_tokens` are **caller-supplied**
/// — `LlmIo` is a passive carrier and the projector trusts the sidecar
/// verbatim.
///
/// In compliance-grade audit-trail use cases (signed evidence, regulator
/// review, billing), these fields MUST originate directly from the
/// `LlmClient`'s response (response headers, parsed usage block, or the
/// transport layer). They MUST NOT be inferred or asserted by the agent
/// layer: a compromised tool or agent that constructs `LlmIo` itself can
/// spoof `provider = "ollama"` (zero rate) to under-report cost, or attribute
/// usage to a cheaper model tier.
///
/// The projector also does not cross-check these fields against
/// `Episode::LlmCall.tokens`. Adding that consistency check (e.g. asserting
/// `prompt_tokens + completion_tokens == LlmCall.tokens`) is a deferred
/// follow-up; for now, callers wanting that guarantee should assert it
/// themselves before passing the sidecar to
/// [`crate::projector::project_with_price_table`].
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct LlmIo {
    /// Recorded prompt text (the `User` message content sent to the LLM).
    pub prompt: String,
    /// Recorded completion text (the assistant response).
    pub completion: String,
    /// Optional model name; populates `Step.name` when present.
    pub model: Option<String>,
    /// Optional provider name (e.g. `"anthropic"`, `"openai"`). Required for
    /// price-table lookup in [`crate::projector::project_with_price_table`].
    #[serde(default)]
    pub provider: Option<String>,
    /// Optional prompt-token count. Required for price-table cost accounting.
    #[serde(default)]
    pub prompt_tokens: Option<u32>,
    /// Optional completion-token count. Required for price-table cost
    /// accounting.
    #[serde(default)]
    pub completion_tokens: Option<u32>,
}

/// Aggregate token usage across all `LlmCall` steps in a run.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Usage {
    /// Total prompt tokens.
    pub prompt_tokens: u32,
    /// Total completion tokens.
    pub completion_tokens: u32,
}

/// Optional cost estimate for the run, in USD.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Cost {
    /// Prompt cost in USD.
    pub prompt_usd: f64,
    /// Completion cost in USD.
    pub completion_usd: f64,
    /// Total cost in USD.
    pub total_usd: f64,
}

/// The aggregate RunLog view per spec §8.2.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RunLog {
    /// Run identifier.
    #[cfg_attr(feature = "schemars", schemars(with = "String"))]
    pub run_id: RunId,
    /// Agent name.
    pub agent: String,
    /// First-event timestamp.
    pub started_at: DateTime<Utc>,
    /// Last-event timestamp, populated when the run is no longer `Running`.
    pub finished_at: Option<DateTime<Utc>>,
    /// Lifecycle status.
    pub status: RunStatus,
    /// Ordered list of steps.
    pub steps: Vec<Step>,
    /// Aggregate usage.
    pub tokens: Usage,
    /// Optional cost estimate; `None` until a pricing table is wired (Wave 7+ carryover).
    pub cost_estimate: Option<Cost>,
}

mod duration_millis {
    use serde::{Deserialize, Deserializer, Serializer};
    use std::time::Duration;

    pub fn serialize<S: Serializer>(d: &Duration, s: S) -> Result<S::Ok, S::Error> {
        s.serialize_u64(d.as_millis() as u64)
    }
    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Duration, D::Error> {
        let ms = u64::deserialize(d)?;
        Ok(Duration::from_millis(ms))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use klieo_core::ids::RunId;
    use std::time::Duration;

    #[test]
    fn run_status_serialises_snake_case() {
        let json = serde_json::to_string(&RunStatus::Running).unwrap();
        assert_eq!(json, "\"running\"");
    }

    #[test]
    fn step_kind_round_trip() {
        for k in [StepKind::LlmCall, StepKind::ToolCall] {
            let json = serde_json::to_string(&k).unwrap();
            let back: StepKind = serde_json::from_str(&json).unwrap();
            assert_eq!(k, back);
        }
    }

    #[test]
    fn step_round_trip() {
        let step = Step {
            idx: 0,
            kind: StepKind::ToolCall,
            name: Some("calc".into()),
            input: serde_json::json!({"a": 1}),
            output: serde_json::json!({"b": 2}),
            error: None,
            latency: Duration::from_millis(42),
            span_id: Some("span-xyz".into()),
        };
        let json = serde_json::to_string(&step).unwrap();
        let back: Step = serde_json::from_str(&json).unwrap();
        assert_eq!(step.idx, back.idx);
        assert_eq!(step.latency, back.latency);
        assert_eq!(step.name, back.name);
    }

    #[test]
    fn run_log_round_trip_empty_steps() {
        let log = RunLog {
            run_id: RunId::new(),
            agent: "writer".into(),
            started_at: Utc::now(),
            finished_at: None,
            status: RunStatus::Running,
            steps: vec![],
            tokens: Usage::default(),
            cost_estimate: None,
        };
        let json = serde_json::to_string(&log).unwrap();
        assert!(json.contains("\"steps\":[]"));
        let _back: RunLog = serde_json::from_str(&json).unwrap();
    }

    #[test]
    fn usage_default_is_zero() {
        let u = Usage::default();
        assert_eq!(u.prompt_tokens, 0);
        assert_eq!(u.completion_tokens, 0);
    }

    #[test]
    fn cost_round_trip() {
        let c = Cost {
            prompt_usd: 0.001,
            completion_usd: 0.002,
            total_usd: 0.003,
        };
        let json = serde_json::to_string(&c).unwrap();
        let back: Cost = serde_json::from_str(&json).unwrap();
        assert!((back.total_usd - 0.003).abs() < 1e-9);
    }
}
