//! Reduced rollout trace model.
//!
//! These types describe the deterministic replay output. They intentionally
//! separate model-visible conversation from runtime/debug objects.

use std::collections::BTreeMap;

use serde::Deserialize;
use serde::Serialize;

use crate::payload::RawPayloadId;
use crate::payload::RawPayloadRef;
mod conversation;
mod runtime;
mod session;

pub use conversation::*;
pub use runtime::*;
pub use session::*;

/// Codex conversation/session UUID.
pub type AgentThreadId = String;
/// Stable multi-agent routing path such as `/root` or `/root/search_docs`.
pub type AgentPath = String;
/// Runtime submission/activation UUID. This is not a chat turn.
pub type CodexTurnId = String;
/// Reduced transcript item ID assigned by the trace reducer.
pub type ConversationItemId = String;
/// Local ID for one outbound upstream inference request.
pub type InferenceCallId = String;
/// Reducer-owned ID for one runtime tool-call object.
pub type ToolCallId = String;
/// Responses `call_id` / custom-tool call ID visible in inference payloads.
pub type ModelVisibleCallId = String;
/// Tool invocation ID assigned inside the code-mode JavaScript runtime.
pub type CodeModeRuntimeToolId = String;
/// Reducer-owned ID for one model-authored `exec` JavaScript cell.
pub type CodeCellId = String;
/// Process/session ID returned by Codex's terminal runtime.
pub type TerminalId = String;
/// Reducer-owned ID for one command/write/poll operation against a terminal.
pub type TerminalOperationId = String;
/// Reducer-owned ID for one installed conversation-history checkpoint.
pub type CompactionId = String;
/// Reducer-owned ID for one upstream request that computes a compaction.
pub type CompactionRequestId = String;
/// Reducer-owned ID for one information-flow edge.
pub type EdgeId = String;
/// Reducer-owned ID for request/log correlation metadata.
pub type CorrelationId = String;

/// Canonical reduced graph for one Codex rollout.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RolloutTrace {
    pub schema_version: u32,
    /// Unique identity for this trace capture.
    ///
    /// `rollout_id` names the Codex rollout/session being observed. `trace_id`
    /// names the diagnostic artifact produced for that rollout, which keeps
    /// storage/replay identity separate from the product-level session identity.
    pub trace_id: String,
    /// CLI-visible rollout/run identity. Higher-level experiment/sample IDs wrap this object.
    pub rollout_id: String,
    pub started_at_unix_ms: i64,
    /// Wall-clock timestamp for terminal rollout status. `None` means running or partial trace.
    pub ended_at_unix_ms: Option<i64>,
    pub status: RolloutStatus,
    pub root_thread_id: AgentThreadId,
    pub threads: BTreeMap<AgentThreadId, AgentThread>,
    pub codex_turns: BTreeMap<CodexTurnId, CodexTurn>,
    pub conversation_items: BTreeMap<ConversationItemId, ConversationItem>,
    pub inference_calls: BTreeMap<InferenceCallId, InferenceCall>,
    /// Model-authored `exec` JavaScript cells keyed by reducer-owned cell ID.
    pub code_cells: BTreeMap<CodeCellId, CodeCell>,
    pub tool_calls: BTreeMap<ToolCallId, ToolCall>,
    /// Terminal runtime sessions keyed by process/session ID returned by the runtime.
    pub terminal_sessions: BTreeMap<TerminalId, TerminalSession>,
    /// Commands/writes/polls against terminals keyed by reducer-owned operation ID.
    pub terminal_operations: BTreeMap<TerminalOperationId, TerminalOperation>,
    /// Installed compaction checkpoints keyed by checkpoint ID.
    pub compactions: BTreeMap<CompactionId, Compaction>,
    /// Upstream remote compaction calls keyed by local request ID.
    pub compaction_requests: BTreeMap<CompactionRequestId, CompactionRequest>,
    /// Information-flow edges between threads, cells, tools, and runtime resources.
    pub interaction_edges: BTreeMap<EdgeId, InteractionEdge>,
    /// Raw JSON payloads keyed by raw-payload ID. Most point at files outside this object.
    pub raw_payloads: BTreeMap<RawPayloadId, RawPayloadRef>,
}

impl RolloutTrace {
    /// Builds an empty reduced trace that a reducer can populate.
    pub(crate) fn new(
        schema_version: u32,
        trace_id: String,
        rollout_id: String,
        root_thread_id: AgentThreadId,
        started_at_unix_ms: i64,
    ) -> Self {
        Self {
            schema_version,
            trace_id,
            rollout_id,
            started_at_unix_ms,
            ended_at_unix_ms: None,
            status: RolloutStatus::Running,
            root_thread_id,
            threads: BTreeMap::new(),
            codex_turns: BTreeMap::new(),
            conversation_items: BTreeMap::new(),
            inference_calls: BTreeMap::new(),
            code_cells: BTreeMap::new(),
            tool_calls: BTreeMap::new(),
            terminal_sessions: BTreeMap::new(),
            terminal_operations: BTreeMap::new(),
            compactions: BTreeMap::new(),
            compaction_requests: BTreeMap::new(),
            interaction_edges: BTreeMap::new(),
            raw_payloads: BTreeMap::new(),
        }
    }
}