transcript.rs

//! Core types for chat transcript
//!
//! This module contains the types that represent the conversation transcript.
//! A Transcript contains Turns, and each Turn contains Blocks.
#![allow(dead_code)]

use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};

use chrono::{DateTime, Utc};
#[cfg(feature = "cli")]
use ratatui::{
    style::{Color, Modifier, Style},
    text::{Line, Span},
};
use serde::{ser::SerializeStruct, Deserialize, Serialize};

#[cfg(feature = "cli")]
use crate::compaction::CompactionBlock;
use crate::config::{CODEY_DIR, TRANSCRIPTS_DIR};
#[cfg(feature = "cli")]
use crate::tools::io::{format_for_user, DEFAULT_TAB_WIDTH};

/// Global counter for unique block IDs
static NEXT_BLOCK_ID: AtomicUsize = AtomicUsize::new(1);

/// Generate a unique block ID
pub fn next_block_id() -> usize {
    NEXT_BLOCK_ID.fetch_add(1, Ordering::Relaxed)
}

/// Role of the message sender
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    User,
    Assistant,
    System,
}

/// Status of a message, tool, or action
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Status {
    Pending,
    Running,
    Complete,
    Error,
    Denied,
    Cancelled,
}

/// Block type identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockType {
    Text,
    Thinking,
    Tool,
    Compaction,
}

/// Get the transcripts directory path, creating it if necessary
fn get_transcripts_dir() -> std::io::Result<PathBuf> {
    let dir = PathBuf::from(CODEY_DIR).join(TRANSCRIPTS_DIR);
    if !dir.exists() {
        std::fs::create_dir_all(&dir)?;
    }
    Ok(dir)
}

/// Find the latest transcript number by scanning the transcripts directory
fn find_latest_transcript_number(dir: &Path) -> Option<u32> {
    std::fs::read_dir(dir)
        .ok()?
        .filter_map(|entry| entry.ok())
        .filter_map(|entry| {
            let name = entry.file_name();
            let name = name.to_str()?;
            if name.ends_with(".json") {
                name.trim_end_matches(".json").parse::<u32>().ok()
            } else {
                None
            }
        })
        .max()
}

/// Get the path for a transcript with a given number
fn transcript_path(dir: &Path, number: u32) -> PathBuf {
    dir.join(format!("{:06}.json", number))
}

/// Trait for all blocks in a turn
#[typetag::serde(tag = "type")]
pub trait Block: Send + Sync {
    /// Get the kind of this block
    fn kind(&self) -> BlockType;

    /// Render this block to terminal lines with given width for wrapping (CLI only)
    #[cfg(feature = "cli")]
    fn render(&self, width: u16) -> Vec<Line<'_>>;

    /// Get the block's unique ID (0 if unassigned)
    fn id(&self) -> usize {
        0
    }

    /// Set the block's ID (called by container when adding)
    fn set_id(&mut self, _id: usize) {}

    /// Get the status of this block
    fn status(&self) -> Status;

    /// Set the status of this block
    fn set_status(&mut self, status: Status);

    /// Whether this block is ephemeral (rendered but not persisted)
    /// Ephemeral blocks are filtered out during serialization.
    fn is_ephemeral(&self) -> bool {
        false
    }

    /// Render status icon with appropriate color (CLI only)
    #[cfg(feature = "cli")]
    fn render_status(&self) -> Span<'static> {
        let (icon, color) = match self.status() {
            Status::Pending => ("? ", Color::Yellow),
            Status::Running => ("⚙ ", Color::Blue),
            Status::Complete => ("✓ ", Color::Green),
            Status::Error => ("✗ ", Color::Red),
            Status::Denied => ("⊘ ", Color::DarkGray),
            Status::Cancelled => ("⊘ ", Color::Yellow),
        };
        Span::styled(icon, Style::default().fg(color))
    }

    /// Append text content to this block (for streaming)
    fn append_text(&mut self, _text: &str) {}

    /// Get the text content of this block (for restoring agent context)
    fn text(&self) -> Option<&str> {
        None
    }

    /// Get the tool call ID (for restoring agent context)
    fn call_id(&self) -> Option<&str> {
        None
    }

    /// Get the tool name (for restoring agent context)
    fn tool_name(&self) -> Option<&str> {
        None
    }

    /// Get the tool params (for restoring agent context)
    fn params(&self) -> Option<&serde_json::Value> {
        None
    }

    /// Set the agent label (for sub-agent tools)
    fn set_agent_label(&mut self, _label: String) {}

    /// Get the agent label (for sub-agent tools)
    fn agent_label(&self) -> Option<&str> {
        None
    }
}

/// Macro to implement common Block trait methods for blocks with text and status fields
#[macro_export]
macro_rules! impl_base_block {
    ($block_type:expr) => {
        fn kind(&self) -> BlockType {
            $block_type
        }

        fn id(&self) -> usize {
            self.id
        }

        fn set_id(&mut self, id: usize) {
            self.id = id;
        }

        fn status(&self) -> Status {
            self.status
        }

        fn set_status(&mut self, status: Status) {
            self.status = status;
        }

        fn append_text(&mut self, text: &str) {
            self.text.push_str(text);
        }

        fn text(&self) -> Option<&str> {
            Some(&self.text)
        }
    };
}

/// Macro for tool blocks that don't have their own ID field.
/// These blocks use call_id for identification instead.
#[macro_export]
macro_rules! impl_tool_block {
    ($block_type:expr) => {
        fn kind(&self) -> BlockType {
            $block_type
        }

        fn status(&self) -> Status {
            self.status
        }

        fn set_status(&mut self, status: Status) {
            self.status = status;
        }

        fn append_text(&mut self, text: &str) {
            self.text.push_str(text);
        }

        fn text(&self) -> Option<&str> {
            Some(&self.text)
        }
    };
}

/// Macro to define a complete tool block with minimal boilerplate.
///
/// This macro generates the struct definition, constructors, and Block trait
/// implementation. You only need to provide the custom header rendering logic.
///
/// # Usage
///
/// ```ignore
/// define_tool_block! {
///     /// Doc comment for the block
///     pub struct MyToolBlock {
///         // Maximum lines to show in result output (default: 10)
///         max_lines: 5,
///         // Type to validate params against
///         params_type: MyToolParams,
///         // Custom header rendering - returns Vec<Span>
///         render_header(self, params) {
///             vec![
///                 Span::styled("my_tool", Style::default().fg(Color::Magenta)),
///                 Span::styled("(", Style::default().fg(Color::DarkGray)),
///                 Span::styled(params["arg"].as_str().unwrap_or(""), Style::default().fg(Color::Cyan)),
///                 Span::styled(")", Style::default().fg(Color::DarkGray)),
///             ]
///         }
///     }
/// }
/// ```
#[macro_export]
macro_rules! define_tool_block {
    (
        $(#[$attr:meta])*
        $vis:vis struct $name:ident {
            max_lines: $max_lines:expr,
            params_type: $params_type:ty,
            render_header($self:ident, $params:ident) $render_body:block
        }
    ) => {
        $(#[$attr])*
        #[derive(Debug, Clone, Serialize, Deserialize)]
        $vis struct $name {
            pub call_id: String,
            pub tool_name: String,
            pub params: serde_json::Value,
            pub status: Status,
            pub text: String,
            #[serde(default)]
            pub background: bool,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub agent_label: Option<String>,
        }

        impl $name {
            pub fn new(
                call_id: impl Into<String>,
                tool_name: impl Into<String>,
                params: serde_json::Value,
                background: bool,
            ) -> Self {
                Self {
                    call_id: call_id.into(),
                    tool_name: tool_name.into(),
                    params,
                    status: Status::Pending,
                    text: String::new(),
                    background,
                    agent_label: None,
                }
            }

            pub fn from_params(
                call_id: &str,
                tool_name: &str,
                params: serde_json::Value,
                background: bool,
            ) -> Option<Self> {
                let _: $params_type = serde_json::from_value(params.clone()).ok()?;
                Some(Self::new(call_id, tool_name, params, background))
            }

            #[allow(unused_variables)]
            fn render_header_spans(&$self) -> Vec<Span<'static>> {
                let $params = &$self.params;
                $render_body
            }
        }

        #[typetag::serde]
        impl Block for $name {
            $crate::impl_tool_block!(BlockType::Tool);

            fn render(&self, _width: u16) -> Vec<Line<'_>> {
                let mut lines = Vec::new();

                // Build header line: status + agent_label + prefix + custom spans
                let mut spans = vec![
                    self.render_status(),
                    render_agent_label(self.agent_label.as_deref()),
                    render_prefix(self.background),
                ];
                spans.extend(self.render_header_spans());
                lines.push(Line::from(spans));

                // Approval prompt if pending
                if self.status == Status::Pending {
                    lines.push(render_approval_prompt());
                }

                // Result output
                if !self.text.is_empty() {
                    lines.extend(render_result(&self.text, $max_lines));
                }

                // Denied message
                if self.status == Status::Denied {
                    lines.push(Line::from(Span::styled(
                        "  Denied by user",
                        Style::default().fg(Color::DarkGray),
                    )));
                }

                lines
            }

            fn call_id(&self) -> Option<&str> {
                Some(&self.call_id)
            }

            fn tool_name(&self) -> Option<&str> {
                Some(&self.tool_name)
            }

            fn params(&self) -> Option<&serde_json::Value> {
                Some(&self.params)
            }

            fn set_agent_label(&mut self, label: String) {
                self.agent_label = Some(label);
            }

            fn agent_label(&self) -> Option<&str> {
                self.agent_label.as_deref()
            }
        }
    };
}

/// Macro for simple tool blocks without agent_label support.
/// Used for internal tools like list_agents, get_agent, list_background_tasks, etc.
#[macro_export]
macro_rules! define_simple_tool_block {
    (
        $(#[$attr:meta])*
        $vis:vis struct $name:ident {
            max_lines: $max_lines:expr,
            render_header($self:ident, $params:ident) $render_body:block
        }
    ) => {
        $(#[$attr])*
        #[derive(Debug, Clone, Serialize, Deserialize)]
        $vis struct $name {
            pub call_id: String,
            pub tool_name: String,
            pub params: serde_json::Value,
            pub status: Status,
            pub text: String,
            #[serde(default)]
            pub background: bool,
        }

        impl $name {
            pub fn new(
                call_id: impl Into<String>,
                tool_name: impl Into<String>,
                params: serde_json::Value,
                background: bool,
            ) -> Self {
                Self {
                    call_id: call_id.into(),
                    tool_name: tool_name.into(),
                    params,
                    status: Status::Pending,
                    text: String::new(),
                    background,
                }
            }

            #[allow(unused_variables)]
            fn render_header_spans(&$self) -> Vec<Span<'static>> {
                let $params = &$self.params;
                $render_body
            }
        }

        #[typetag::serde]
        impl Block for $name {
            $crate::impl_tool_block!(BlockType::Tool);

            fn render(&self, _width: u16) -> Vec<Line<'_>> {
                let mut lines = Vec::new();

                // Build header line: status + prefix + custom spans
                let mut spans = vec![
                    self.render_status(),
                    render_prefix(self.background),
                ];
                spans.extend(self.render_header_spans());
                lines.push(Line::from(spans));

                // Approval prompt if pending
                if self.status == Status::Pending {
                    lines.push(render_approval_prompt());
                }

                // Result output
                if !self.text.is_empty() {
                    lines.extend(render_result(&self.text, $max_lines));
                }

                // Denied message
                if self.status == Status::Denied {
                    lines.push(Line::from(Span::styled(
                        "  Denied by user",
                        Style::default().fg(Color::DarkGray),
                    )));
                }

                lines
            }

            fn call_id(&self) -> Option<&str> {
                Some(&self.call_id)
            }

            fn tool_name(&self) -> Option<&str> {
                Some(&self.tool_name)
            }

            fn params(&self) -> Option<&serde_json::Value> {
                Some(&self.params)
            }
        }
    };
}

/// Simple text content
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TextBlock {
    #[serde(default = "next_block_id")]
    pub id: usize,
    pub text: String,
    pub status: Status,
}

impl TextBlock {
    // TODO Delete `new` and force using a keyword to create
    pub fn new(text: impl Into<String>) -> Self {
        Self {
            id: next_block_id(),
            text: text.into(),
            status: Status::Running,
        }
    }

    pub fn pending(text: impl Into<String>) -> Self {
        Self {
            id: next_block_id(),
            text: text.into(),
            status: Status::Pending,
        }
    }

    pub fn complete(text: impl Into<String>) -> Self {
        Self {
            id: next_block_id(),
            text: text.into(),
            status: Status::Complete,
        }
    }
}

#[typetag::serde]
impl Block for TextBlock {
    impl_base_block!(BlockType::Text);

    #[cfg(feature = "cli")]
    fn render(&self, width: u16) -> Vec<Line<'_>> {
        // Use ratskin for markdown rendering
        let skin = ratskin::RatSkin::default();
        let text = ratskin::RatSkin::parse_text(&self.text);
        skin.parse(text, width)
    }
}

/// Thinking/reasoning content (extended thinking)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThinkingBlock {
    #[serde(default = "next_block_id")]
    pub id: usize,
    pub text: String,
    pub status: Status,
}

impl ThinkingBlock {
    pub fn new(text: impl Into<String>) -> Self {
        Self {
            id: next_block_id(),
            text: text.into(),
            status: Status::Running,
        }
    }
}

#[typetag::serde]
impl Block for ThinkingBlock {
    impl_base_block!(BlockType::Thinking);

    #[cfg(feature = "cli")]
    fn render(&self, width: u16) -> Vec<Line<'_>> {
        let mut lines = Vec::new();
        let style = Style::default()
            .fg(Color::DarkGray)
            .add_modifier(Modifier::ITALIC);

        let wrapped = textwrap::wrap(&self.text, width as usize);
        for line in wrapped {
            lines.push(Line::from(Span::styled(line, style)));
        }
        lines
    }
}

/// Generic tool content (fallback for tools without specialized display)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolBlock {
    #[serde(default = "next_block_id")]
    pub id: usize,
    pub call_id: String,
    pub name: String,
    pub params: serde_json::Value,
    pub status: Status,
    pub text: String,
    /// If true, tool runs in background
    #[serde(default)]
    pub background: bool,
    /// Agent label for sub-agent tools
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub agent_label: Option<String>,
}

impl ToolBlock {
    pub fn new(
        call_id: impl Into<String>,
        name: impl Into<String>,
        params: serde_json::Value,
        background: bool,
    ) -> Self {
        Self {
            id: next_block_id(),
            call_id: call_id.into(),
            name: name.into(),
            params,
            status: Status::Pending,
            text: String::new(),
            background,
            agent_label: None,
        }
    }
}

#[typetag::serde]
impl Block for ToolBlock {
    impl_base_block!(BlockType::Tool);

    #[cfg(feature = "cli")]
    fn render(&self, _width: u16) -> Vec<Line<'_>> {
        let mut lines = Vec::new();

        // Tool name with status icon, optional agent label, and optional [bg] prefix
        lines.push(Line::from(vec![
            self.render_status(),
            render_agent_label(self.agent_label.as_deref()),
            render_prefix(self.background),
            Span::styled(
                &self.name,
                Style::default()
                    .fg(Color::Magenta)
                    .add_modifier(Modifier::BOLD),
            ),
        ]));

        // Params
        let params_str = serde_json::to_string_pretty(&self.params).unwrap_or_default();
        for param_line in params_str.lines().take(10) {
            lines.push(Line::from(Span::styled(
                format!("  {}", param_line),
                Style::default().fg(Color::DarkGray),
            )));
        }
        if params_str.lines().count() > 10 {
            lines.push(Line::from(Span::styled(
                "  ...",
                Style::default().fg(Color::DarkGray),
            )));
        }

        // Approval prompt if pending
        if self.status == Status::Pending {
            lines.push(render_approval_prompt());
        }

        // Result if completed
        if !self.text.is_empty() {
            lines.extend(render_result(&self.text, 5));
        }

        // Denied message
        if self.status == Status::Denied {
            lines.push(Line::from(Span::styled(
                "  Denied by user",
                Style::default().fg(Color::DarkGray),
            )));
        }

        lines
    }

    fn call_id(&self) -> Option<&str> {
        Some(&self.call_id)
    }

    fn tool_name(&self) -> Option<&str> {
        Some(&self.name)
    }

    fn params(&self) -> Option<&serde_json::Value> {
        Some(&self.params)
    }

    fn set_agent_label(&mut self, label: String) {
        self.agent_label = Some(label);
    }

    fn agent_label(&self) -> Option<&str> {
        self.agent_label.as_deref()
    }
}

/// Notification block for mid-turn injected messages
/// These are ephemeral - rendered but not persisted to the transcript.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NotificationBlock {
    #[serde(default = "next_block_id")]
    pub id: usize,
    pub source: String,
    pub text: String,
    #[serde(skip_deserializing, default = "NotificationBlock::default_status")]
    status: Status,
}

impl NotificationBlock {
    pub fn new(source: impl Into<String>, text: impl Into<String>) -> Self {
        Self {
            id: next_block_id(),
            source: source.into(),
            text: text.into(),
            status: Status::Complete,
        }
    }

    fn default_status() -> Status {
        Status::Complete
    }
}

#[typetag::serde]
impl Block for NotificationBlock {
    impl_base_block!(BlockType::Text);

    fn is_ephemeral(&self) -> bool {
        true // Key difference - not persisted
    }

    #[cfg(feature = "cli")]
    fn render(&self, _width: u16) -> Vec<Line<'_>> {
        // Single line: » source: text
        vec![Line::from(vec![
            Span::styled("» ", Style::default().fg(Color::Yellow)),
            Span::styled(&self.text, Style::default().fg(Color::Yellow)),
        ])]
    }
}

/// Helper: render prefix for background tools - "[bg] " if true, empty otherwise
#[cfg(feature = "cli")]
pub fn render_prefix(background: bool) -> Span<'static> {
    if background {
        Span::styled("[bg] ", Style::default().fg(Color::Cyan))
    } else {
        Span::raw("")
    }
}

/// Helper: render agent label prefix for sub-agent tools
#[cfg(feature = "cli")]
pub fn render_agent_label(label: Option<&str>) -> Span<'static> {
    match label {
        Some(l) => Span::styled(format!("[{}] ", l), Style::default().fg(Color::Cyan)),
        None => Span::raw(""),
    }
}

/// Helper: render approval prompt
#[cfg(feature = "cli")]
pub fn render_approval_prompt() -> Line<'static> {
    Line::from(vec![
        Span::styled("  [", Style::default().fg(Color::DarkGray)),
        Span::styled(
            "y",
            Style::default()
                .fg(Color::Green)
                .add_modifier(Modifier::BOLD),
        ),
        Span::styled("]es  [", Style::default().fg(Color::DarkGray)),
        Span::styled(
            "n",
            Style::default().fg(Color::Red).add_modifier(Modifier::BOLD),
        ),
        Span::styled("]o", Style::default().fg(Color::DarkGray)),
    ])
}

/// Helper: render result with line limit
#[cfg(feature = "cli")]
pub fn render_result(result: &str, max_lines: usize) -> Vec<Line<'static>> {
    let mut lines = Vec::new();
    let preview_lines: Vec<&str> = result.lines().take(max_lines).collect();
    for line in &preview_lines {
        let formatted = format_for_user(line, DEFAULT_TAB_WIDTH);
        lines.push(Line::from(Span::styled(
            format!("  {}", formatted),
            Style::default().fg(Color::DarkGray),
        )));
    }
    if result.lines().count() > max_lines {
        lines.push(Line::from(Span::styled(
            "  ...",
            Style::default().fg(Color::DarkGray),
        )));
    }
    lines
}

/// Staging area for pending blocks (notifications, approvals) awaiting consumption.
/// Similar to Turn but without role/timestamp - blocks here haven't entered the conversation yet.
#[derive(Default)]
pub struct Stage {
    blocks: Vec<Box<dyn Block>>,
}

impl Stage {
    pub fn new() -> Self {
        Self { blocks: Vec::new() }
    }

    /// Add a block to staging, returns its ID
    pub fn push(&mut self, block: Box<dyn Block>) -> usize {
        let id = block.id();
        self.blocks.push(block);
        id
    }

    /// Get a block by ID
    pub fn get(&self, id: usize) -> Option<&Box<dyn Block>> {
        self.blocks.iter().find(|b| b.id() == id)
    }

    /// Get a mutable block by ID
    pub fn get_mut(&mut self, id: usize) -> Option<&mut Box<dyn Block>> {
        self.blocks.iter_mut().find(|b| b.id() == id)
    }

    /// Remove and return a block by ID (for promotion to transcript)
    pub fn remove(&mut self, id: usize) -> Option<Box<dyn Block>> {
        if let Some(pos) = self.blocks.iter().position(|b| b.id() == id) {
            Some(self.blocks.remove(pos))
        } else {
            None
        }
    }

    /// Check if staging is empty
    pub fn is_empty(&self) -> bool {
        self.blocks.is_empty()
    }

    /// Number of staged blocks
    pub fn len(&self) -> usize {
        self.blocks.len()
    }

    /// Iterate over blocks (for rendering)
    pub fn iter(&self) -> impl Iterator<Item = &Box<dyn Block>> {
        self.blocks.iter()
    }

    /// Drain all blocks (for batch promotion)
    pub fn drain_all(&mut self) -> Vec<Box<dyn Block>> {
        self.blocks.drain(..).collect()
    }

    /// Find a tool block by call_id and return mutable reference
    pub fn find_by_call_id(&mut self, call_id: &str) -> Option<&mut Box<dyn Block>> {
        self.blocks.iter_mut().find(|b| b.call_id() == Some(call_id))
    }

    /// Remove a tool block by call_id (for promotion to transcript)
    pub fn remove_by_call_id(&mut self, call_id: &str) -> Option<Box<dyn Block>> {
        if let Some(pos) = self.blocks.iter().position(|b| b.call_id() == Some(call_id)) {
            Some(self.blocks.remove(pos))
        } else {
            None
        }
    }

    /// Render all staged blocks with given width (CLI only)
    #[cfg(feature = "cli")]
    pub fn render(&self, width: u16) -> Vec<Line<'_>> {
        let mut lines = Vec::new();
        for (i, block) in self.blocks.iter().enumerate() {
            lines.extend(block.render(width));
            // Add blank line between blocks (but not after last)
            if i < self.blocks.len() - 1 {
                lines.push(Line::from(""));
            }
        }
        lines
    }
}

/// A turn in the conversation - one user or assistant response
#[derive(Deserialize)]
pub struct Turn {
    pub id: usize,
    pub role: Role,
    pub content: Vec<Box<dyn Block>>,
    pub timestamp: DateTime<Utc>,
    /// Index of the currently active (streaming) block, if any
    #[serde(skip)]
    pub active_block_idx: Option<usize>,
}

/// Custom serialization that filters out ephemeral blocks
impl Serialize for Turn {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        // Filter out ephemeral blocks before serializing
        let persistent_content: Vec<&Box<dyn Block>> =
            self.content.iter().filter(|b| !b.is_ephemeral()).collect();

        let mut state = serializer.serialize_struct("Turn", 4)?;
        state.serialize_field("id", &self.id)?;
        state.serialize_field("role", &self.role)?;
        state.serialize_field("content", &persistent_content)?;
        state.serialize_field("timestamp", &self.timestamp)?;
        state.end()
    }
}

impl Turn {
    pub fn new(id: usize, role: Role, content: Vec<Box<dyn Block>>) -> Self {
        Self {
            id,
            role,
            content,
            timestamp: Utc::now(),
            active_block_idx: None,
        }
    }

    /// Add a block and return its index
    pub fn add_block(&mut self, block: Box<dyn Block>) -> usize {
        let idx = self.content.len();
        self.content.push(block);
        idx
    }

    /// Append text to a specific block by index
    pub fn append_to_block(&mut self, idx: usize, text: &str) {
        if let Some(block) = self.content.get_mut(idx) {
            block.append_text(text);
        }
    }

    /// Get a mutable block by index
    pub fn get_block_mut(&mut self, idx: usize) -> Option<&mut (dyn Block + 'static)> {
        self.content.get_mut(idx).map(|b| b.as_mut())
    }

    /// Mark a block as complete by index
    pub fn complete_block(&mut self, idx: usize) {
        if let Some(block) = self.get_block_mut(idx) {
            block.set_status(Status::Complete);
        }
    }

    /// Check if the active block matches the expected type
    pub fn is_active_block_type(&self, expected: BlockType) -> bool {
        self.active_block_idx
            .and_then(|idx| self.content.get(idx))
            .map(|block| block.kind() == expected)
            .unwrap_or(false)
    }

    /// Start a new block (completes previous active block if any)
    pub fn start_block(&mut self, block: Box<dyn Block>) -> usize {
        if let Some(prev_idx) = self.active_block_idx {
            self.complete_block(prev_idx);
        }
        let idx = self.add_block(block);
        self.active_block_idx = Some(idx);
        idx
    }

    /// Append text to the currently active block
    pub fn append_to_active(&mut self, text: &str) {
        if let Some(idx) = self.active_block_idx {
            self.append_to_block(idx, text);
        }
    }

    /// Get a mutable reference to the active block
    pub fn get_active_block_mut(&mut self) -> Option<&mut (dyn Block + 'static)> {
        self.active_block_idx
            .and_then(|idx| self.get_block_mut(idx))
    }

    /// Render all blocks with given width (CLI only)
    #[cfg(feature = "cli")]
    pub fn render(&self, width: u16) -> Vec<Line<'_>> {
        let mut lines = Vec::new();
        for (i, block) in self.content.iter().enumerate() {
            lines.extend(block.render(width));
            // Add blank line between blocks (but not after last)
            if i < self.content.len() - 1 {
                lines.push(Line::from(""));
            }
        }
        lines
    }
}

/// The chat transcript - display log of all turns for UI rendering
#[derive(Serialize, Deserialize)]
pub struct Transcript {
    turns: Vec<Turn>,
    next_id: usize,
    #[serde(skip)]
    path: Option<PathBuf>,
    /// ID of the current turn being streamed to (if any)
    #[serde(skip)]
    current_turn_id: Option<usize>,
    /// Staging area for pending blocks
    #[serde(skip)]
    pub stage: Stage,
}

impl Transcript {
    /// Create a new transcript with a specific path
    pub fn with_path(path: PathBuf) -> Self {
        Self {
            turns: Vec::new(),
            next_id: 0,
            path: Some(path),
            current_turn_id: None,
            stage: Stage::new(),
        }
    }

    /// Get the current path
    pub fn path(&self) -> Option<&Path> {
        self.path.as_deref()
    }

    fn next_id(&mut self) -> usize {
        let id = self.next_id;
        self.next_id += 1;
        id
    }

    /// Add a new turn with a single block
    pub fn add_turn(&mut self, role: Role, block: impl Block + 'static) -> usize {
        let id = self.next_id();
        self.turns.push(Turn::new(id, role, vec![Box::new(block)]));
        id
    }

    /// Add an empty turn (for streaming)
    pub fn add_empty(&mut self, role: Role) -> usize {
        let id = self.next_id();
        self.turns.push(Turn::new(id, role, vec![]));
        id
    }

    pub fn get_mut(&mut self, id: usize) -> Option<&mut Turn> {
        self.turns.iter_mut().find(|t| t.id == id)
    }

    pub fn turns(&self) -> &[Turn] {
        &self.turns
    }

    // =========================================================================
    // Turn streaming lifecycle
    // =========================================================================

    /// Begin a new turn for streaming. Must call finish_turn() when done.
    pub fn begin_turn(&mut self, role: Role) {
        if self.current_turn_id.is_some() {
            panic!("Cannot begin turn: previous turn not finished");
        }
        let id = self.add_empty(role);
        self.current_turn_id = Some(id);
    }

    /// Finish the current turn - marks active block complete, clears current turn.
    pub fn finish_turn(&mut self) {
        self.mark_active_block(Status::Complete);
        self.current_turn_id = None;
    }

    /// Get mutable reference to the current turn. Panics if no turn is active.
    fn current_turn_mut(&mut self) -> &mut Turn {
        let turn_id = self
            .current_turn_id
            .expect("No active turn - call begin_turn() first");
        self.get_mut(turn_id).expect("Current turn ID is invalid")
    }

    /// Stream a delta to the current turn.
    /// Appends to active block if type matches, otherwise starts a new block.
    /// Panics if no turn is active.
    pub fn stream_delta(&mut self, kind: BlockType, text: &str) {
        let turn = self.current_turn_mut();
        if turn.is_active_block_type(kind) {
            turn.append_to_active(text);
        } else {
            let block: Box<dyn Block> = match kind {
                BlockType::Text => Box::new(TextBlock::new(text)),
                BlockType::Thinking => Box::new(ThinkingBlock::new(text)),
                #[cfg(feature = "cli")]
                BlockType::Compaction => Box::new(CompactionBlock::new(text)),
                #[cfg(not(feature = "cli"))]
                BlockType::Compaction => Box::new(TextBlock::new(text)),
                BlockType::Tool => panic!("Use start_block for tools"),
            };
            turn.start_block(block);
        }
    }

    /// Start a new block on the current turn. Panics if no turn is active.
    pub fn start_block(&mut self, block: Box<dyn Block>) {
        self.current_turn_mut().start_block(block);
    }

    /// Get mutable reference to the active block.
    pub fn active_block_mut(&mut self) -> Option<&mut (dyn Block + 'static)> {
        let turn_id = self.current_turn_id?;
        self.get_mut(turn_id)?.get_active_block_mut()
    }

    /// Find a tool block by its call_id.
    pub fn find_tool_block_mut(&mut self, call_id: &str) -> Option<&mut (dyn Block + 'static)> {
        for turn in &mut self.turns {
            for block in &mut turn.content {
                if block.call_id() == Some(call_id) {
                    return Some(block.as_mut());
                }
            }
        }
        None
    }

    /// Set status on the active block.
    pub fn mark_active_block(&mut self, status: Status) {
        if let Some(block) = self.active_block_mut() {
            block.set_status(status);
        }
    }

    /// Check if current turn's active block matches a type.
    pub fn is_streaming_block_type(&self, kind: BlockType) -> bool {
        self.current_turn_id
            .and_then(|id| self.turns.iter().find(|t| t.id == id))
            .map(|t| t.is_active_block_type(kind))
            .unwrap_or(false)
    }

    // =========================================================================
    // Persistence
    // =========================================================================

    /// Save transcript to its path
    pub fn save(&self) -> std::io::Result<()> {
        let path = self.path.as_ref().ok_or_else(|| {
            std::io::Error::new(std::io::ErrorKind::Other, "No path set for transcript")
        })?;

        // Ensure parent directory exists
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let file = std::fs::File::create(path)?;
        serde_json::to_writer_pretty(file, self)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
    }

    /// Load the latest transcript from the transcripts directory
    /// If no transcripts exist, creates a new one with number 0
    pub fn load() -> std::io::Result<Self> {
        let dir = get_transcripts_dir()?;

        if let Some(latest_number) = find_latest_transcript_number(&dir) {
            let path = transcript_path(&dir, latest_number);
            let file = std::fs::File::open(&path)?;
            let mut transcript: Self = serde_json::from_reader(file)
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
            transcript.path = Some(path);
            Ok(transcript)
        } else {
            // No transcripts exist, create a new one with number 0
            let path = transcript_path(&dir, 0);
            Ok(Self::with_path(path))
        }
    }

    /// Create a new empty transcript with the next available number
    pub fn new_numbered() -> std::io::Result<Self> {
        let dir = get_transcripts_dir()?;

        let next_number = find_latest_transcript_number(&dir)
            .map(|n| n + 1)
            .unwrap_or(0);

        let path = transcript_path(&dir, next_number);
        Ok(Self::with_path(path))
    }

    /// Rotate to a new transcript file
    /// Saves the current transcript and returns a new one with the next numbered path
    /// If the last turn contains a CompactionBlock, it will be added to the new transcript
    pub fn rotate(&self) -> std::io::Result<Self> {
        // Save current transcript
        self.save()?;

        // Get transcripts directory
        let dir = get_transcripts_dir()?;

        // Find next number
        let next_number = find_latest_transcript_number(&dir)
            .map(|n| n + 1)
            .unwrap_or(0);

        // Create new transcript with next path
        let new_path = transcript_path(&dir, next_number);
        let mut new_transcript = Self::with_path(new_path);

        // Check if last turn has a CompactionBlock and carry it over
        #[cfg(feature = "cli")]
        if let Some(last_turn) = self.turns.last() {
            for block in &last_turn.content {
                if block.kind() == BlockType::Compaction {
                    if let Some(summary_text) = block.text() {
                        use crate::compaction::CompactionBlock;
                        let mut compaction_block = CompactionBlock::new(summary_text.to_string());
                        compaction_block.status = Status::Complete;
                        new_transcript.add_turn(Role::Assistant, compaction_block);
                        break; // Only add the first compaction block found
                    }
                }
            }
        }

        Ok(new_transcript)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_text_block_render() {
        let block = TextBlock::new("Hello\nWorld");
        let lines = block.render(80);
        assert!(!lines.is_empty());
    }

    #[test]
    fn test_tool_block_status() {
        let mut block = ToolBlock::new("call_1", "test", serde_json::json!({}), false);
        assert_eq!(block.status(), Status::Pending);

        block.set_status(Status::Running);
        assert_eq!(block.status(), Status::Running);

        block.set_status(Status::Complete);
        block.append_text("done");
        assert_eq!(block.status(), Status::Complete);
    }

    #[test]
    fn test_transcript_add_and_get() {
        let mut transcript = Transcript::with_path(std::path::PathBuf::from("/tmp/test.md"));

        let id1 = transcript.add_turn(Role::User, TextBlock::new("Hello"));
        let id2 = transcript.add_turn(Role::Assistant, TextBlock::new("Hi there!"));

        assert_eq!(transcript.turns().len(), 2);
        assert_eq!(transcript.get_mut(id1).unwrap().role, Role::User);
        assert_eq!(transcript.get_mut(id2).unwrap().role, Role::Assistant);
    }

    #[test]
    fn test_turn_streaming() {
        let mut turn = Turn::new(0, Role::Assistant, vec![]);

        // Start streaming - add a text block and get its index
        let idx = turn.add_block(Box::new(TextBlock::new("Hello")));
        assert_eq!(idx, 0);

        // Append to that block
        turn.append_to_block(idx, " world");

        assert_eq!(turn.content.len(), 1);
    }

    #[test]
    fn test_transcript_save_load_roundtrip() {
        let temp_dir = std::env::temp_dir();
        let path = temp_dir.join("codey_test_transcript.json");

        let mut transcript = Transcript::with_path(path.clone());
        transcript.add_turn(Role::User, TextBlock::new("Hello"));
        transcript.add_turn(Role::Assistant, TextBlock::new("Hi there!"));
        transcript.add_turn(Role::User, TextBlock::new("How are you?"));

        transcript.save().expect("Failed to save transcript");

        // Load by reading the file directly
        let file = std::fs::File::open(&path).expect("Failed to open file");
        let loaded: Transcript = serde_json::from_reader(file).expect("Failed to deserialize");

        assert_eq!(loaded.turns().len(), 3);
        assert_eq!(loaded.turns()[0].role, Role::User);
        assert_eq!(loaded.turns()[1].role, Role::Assistant);
        assert_eq!(loaded.turns()[2].role, Role::User);

        // Verify text content is preserved
        let text = loaded.turns()[0].content[0].text();
        assert_eq!(text, Some("Hello"));

        // Clean up
        let _ = std::fs::remove_file(path);
    }

    #[test]
    fn test_transcript_save_load_with_tool_blocks() {
        let temp_dir = std::env::temp_dir();
        let path = temp_dir.join("codey_test_tool_transcript.json");

        let mut transcript = Transcript::with_path(path.clone());
        transcript.add_turn(Role::User, TextBlock::new("Run ls"));

        // Add an assistant turn with a tool block
        let mut tool_block = ToolBlock::new(
            "call_123",
            "shell",
            serde_json::json!({"command": "ls"}),
            false,
        );
        tool_block.set_status(Status::Complete);
        tool_block.append_text("file1.txt\nfile2.txt");
        transcript.add_turn(Role::Assistant, tool_block);

        transcript.save().expect("Failed to save transcript");

        // Load by reading the file directly
        let file = std::fs::File::open(&path).expect("Failed to open file");
        let loaded: Transcript = serde_json::from_reader(file).expect("Failed to deserialize");

        assert_eq!(loaded.turns().len(), 2);

        // Verify tool block is preserved
        let tool_turn = &loaded.turns()[1];
        assert_eq!(tool_turn.role, Role::Assistant);
        let block = &tool_turn.content[0];
        assert_eq!(block.tool_name(), Some("shell"));
        assert_eq!(block.call_id(), Some("call_123"));
        assert!(block.params().is_some());
        assert_eq!(block.text(), Some("file1.txt\nfile2.txt"));

        // Clean up
        let _ = std::fs::remove_file(path);
    }

    // ========================================================================
    // Stage Tests - Simulating concurrent agent behavior
    // ========================================================================

    #[test]
    fn test_stage_push_and_find_by_call_id() {
        let mut stage = Stage::new();

        // Add multiple tool blocks with different call_ids
        let block1 = ToolBlock::new("call_1", "shell", serde_json::json!({"cmd": "ls"}), false);
        let block2 = ToolBlock::new("call_2", "read_file", serde_json::json!({"path": "foo"}), false);
        let block3 = ToolBlock::new("call_3", "shell", serde_json::json!({"cmd": "pwd"}), false);

        stage.push(Box::new(block1));
        stage.push(Box::new(block2));
        stage.push(Box::new(block3));

        assert_eq!(stage.len(), 3);

        // Find by call_id should work for all
        assert!(stage.find_by_call_id("call_1").is_some());
        assert!(stage.find_by_call_id("call_2").is_some());
        assert!(stage.find_by_call_id("call_3").is_some());
        assert!(stage.find_by_call_id("call_999").is_none());
    }

    #[test]
    fn test_stage_remove_by_call_id_ordering() {
        let mut stage = Stage::new();

        // Simulate rapid queuing from multiple agents
        for i in 0..10 {
            let block = ToolBlock::new(
                &format!("call_{}", i),
                "shell",
                serde_json::json!({"cmd": format!("echo {}", i)}),
                false,
            );
            stage.push(Box::new(block));
        }

        assert_eq!(stage.len(), 10);

        // Remove in non-sequential order (simulating different approval timing)
        let removed = stage.remove_by_call_id("call_5");
        assert!(removed.is_some());
        assert_eq!(removed.unwrap().call_id(), Some("call_5"));
        assert_eq!(stage.len(), 9);

        let removed = stage.remove_by_call_id("call_0");
        assert!(removed.is_some());
        assert_eq!(stage.len(), 8);

        let removed = stage.remove_by_call_id("call_9");
        assert!(removed.is_some());
        assert_eq!(stage.len(), 7);

        // Verify remaining blocks are still findable
        assert!(stage.find_by_call_id("call_1").is_some());
        assert!(stage.find_by_call_id("call_2").is_some());
        assert!(stage.find_by_call_id("call_7").is_some());

        // Removed blocks should not be findable
        assert!(stage.find_by_call_id("call_0").is_none());
        assert!(stage.find_by_call_id("call_5").is_none());
        assert!(stage.find_by_call_id("call_9").is_none());
    }

    #[test]
    fn test_stage_interleaved_push_and_remove() {
        let mut stage = Stage::new();

        // Simulate interleaved operations like what happens with auto-approve
        // Agent 1 queues
        stage.push(Box::new(ToolBlock::new("agent1_call1", "shell", serde_json::json!({}), false)));

        // Agent 2 queues (before agent 1's is processed)
        stage.push(Box::new(ToolBlock::new("agent2_call1", "shell", serde_json::json!({}), false)));

        // Agent 1's call is auto-approved, removed
        let removed = stage.remove_by_call_id("agent1_call1");
        assert!(removed.is_some());

        // Agent 3 queues
        stage.push(Box::new(ToolBlock::new("agent3_call1", "read_file", serde_json::json!({}), false)));

        // Agent 1 queues another
        stage.push(Box::new(ToolBlock::new("agent1_call2", "shell", serde_json::json!({}), false)));

        assert_eq!(stage.len(), 3);

        // Process in arrival order
        let removed = stage.remove_by_call_id("agent2_call1");
        assert!(removed.is_some());
        let removed = stage.remove_by_call_id("agent3_call1");
        assert!(removed.is_some());
        let removed = stage.remove_by_call_id("agent1_call2");
        assert!(removed.is_some());

        assert!(stage.is_empty());
    }

    #[test]
    fn test_stage_double_remove_returns_none() {
        let mut stage = Stage::new();

        stage.push(Box::new(ToolBlock::new("call_1", "shell", serde_json::json!({}), false)));

        // First remove succeeds
        let removed = stage.remove_by_call_id("call_1");
        assert!(removed.is_some());

        // Second remove returns None (important for race condition handling)
        let removed = stage.remove_by_call_id("call_1");
        assert!(removed.is_none());
    }

    #[test]
    fn test_stage_with_transcript_promotion() {
        // Simulate the full flow: Stage -> Transcript
        let mut stage = Stage::new();
        let mut transcript = Transcript::with_path(std::path::PathBuf::from("/tmp/test_promo.md"));

        // Queue multiple approvals
        stage.push(Box::new(ToolBlock::new("call_a", "shell", serde_json::json!({"cmd": "ls"}), false)));
        stage.push(Box::new(ToolBlock::new("call_b", "read_file", serde_json::json!({"path": "x"}), false)));
        stage.push(Box::new(ToolBlock::new("call_c", "write_file", serde_json::json!({"path": "y"}), false)));

        // Start an assistant turn for the transcript
        transcript.begin_turn(Role::Assistant);

        // User approves call_b first (out of order)
        if let Some(mut block) = stage.remove_by_call_id("call_b") {
            block.set_status(Status::Running);
            transcript.start_block(block);
        }

        assert_eq!(stage.len(), 2);

        // Auto-approve call_a
        if let Some(mut block) = stage.remove_by_call_id("call_a") {
            block.set_status(Status::Running);
            transcript.start_block(block);
        }

        // User denies call_c
        if let Some(mut block) = stage.remove_by_call_id("call_c") {
            block.set_status(Status::Denied);
            transcript.start_block(block);
        }

        assert!(stage.is_empty());
        // All blocks added to the same turn
        assert_eq!(transcript.turns().len(), 1);
        assert_eq!(transcript.turns()[0].content.len(), 3);
    }

    #[test]
    fn test_stage_high_volume_stress() {
        let mut stage = Stage::new();

        // Simulate 100 rapid tool calls from multiple agents
        let agents = ["primary", "agent1", "agent2", "agent3"];
        for i in 0..100 {
            let agent = agents[i % agents.len()];
            let block = ToolBlock::new(
                &format!("{}_{}", agent, i),
                if i % 3 == 0 { "shell" } else { "read_file" },
                serde_json::json!({"n": i}),
                i % 5 == 0, // Some are background
            );
            stage.push(Box::new(block));
        }

        assert_eq!(stage.len(), 100);

        // Remove in random-ish order (simulating different approval timings)
        let removal_order: Vec<usize> = vec![
            50, 25, 75, 10, 90, 5, 95, 0, 99, 33, 66, 42, 58, 17, 83,
        ];

        for i in removal_order {
            let agent = agents[i % agents.len()];
            let call_id = format!("{}_{}", agent, i);
            let removed = stage.remove_by_call_id(&call_id);
            assert!(removed.is_some(), "Failed to remove {}", call_id);
        }

        assert_eq!(stage.len(), 85); // 100 - 15 removed

        // Verify double-remove protection
        assert!(stage.remove_by_call_id("primary_50").is_none());
        assert!(stage.remove_by_call_id("agent1_25").is_none());
    }

    #[test]
    fn test_stage_drain_all() {
        let mut stage = Stage::new();

        for i in 0..5 {
            stage.push(Box::new(ToolBlock::new(
                &format!("call_{}", i),
                "shell",
                serde_json::json!({}),
                false,
            )));
        }

        let drained = stage.drain_all();
        assert_eq!(drained.len(), 5);
        assert!(stage.is_empty());

        // Draining again returns empty
        let drained = stage.drain_all();
        assert!(drained.is_empty());
    }

    // ========================================================================
    // Turn lifecycle: sub-agent tool completion must not corrupt turn state
    // ========================================================================

    /// Verifies that begin_turn panics when called with an already-active turn.
    /// This is the invariant that was violated before the fix in handle_tool_event:
    /// sub-agent tool completions with injectable notifications would call
    /// begin_turn() creating an orphaned turn, then the next user submit would
    /// call begin_turn() again and hit this panic.
    #[test]
    #[should_panic(expected = "Cannot begin turn: previous turn not finished")]
    fn test_begin_turn_panics_if_previous_not_finished() {
        let mut transcript = Transcript::with_path(std::path::PathBuf::from("/tmp/test_double_begin.md"));

        transcript.begin_turn(Role::Assistant);
        // Second begin_turn without finishing first -> panic
        transcript.begin_turn(Role::Assistant);
    }

    /// Simulates the correct lifecycle when a sub-agent's tool completes
    /// after the primary agent has finished its turn. The fix in
    /// handle_tool_event guards notification injection behind is_primary,
    /// so sub-agent tool completions skip the finish_turn/begin_turn calls
    /// that would corrupt the transcript turn state.
    #[test]
    fn test_spawn_agent_turn_lifecycle_no_orphan() {
        let mut transcript = Transcript::with_path(std::path::PathBuf::from("/tmp/test_spawn_ok.md"));

        // User sends message
        transcript.add_turn(Role::User, TextBlock::new("Spawn an agent to help me"));

        // Primary agent's assistant turn begins
        transcript.begin_turn(Role::Assistant);
        transcript.stream_delta(BlockType::Text, "I'll spawn an agent.");
        let tool_block = ToolBlock::new(
            "call_spawn_1",
            "mcp_spawn_agent",
            serde_json::json!({"task": "review code"}),
            false,
        );
        transcript.start_block(Box::new(tool_block));

        // Primary agent finishes its turn
        transcript.finish_turn();

        // Sub-agent runs in the background. Its tool completes, but since
        // is_primary is false, the handler does NOT call finish_turn/begin_turn.
        // (No transcript manipulation for sub-agent tool completions.)

        // Later, user submits next message - this should work without panic
        transcript.add_turn(Role::User, TextBlock::new("What did the agent find?"));
        transcript.begin_turn(Role::Assistant);
        transcript.stream_delta(BlockType::Text, "The agent found...");
        transcript.finish_turn();

        // Verify transcript integrity
        assert_eq!(transcript.turns().len(), 4); // user, assistant, user, assistant
    }

    /// Verifies that finish_turn is safe to call when no turn is active.
    #[test]
    fn test_finish_turn_without_active_turn_is_safe() {
        let mut transcript = Transcript::with_path(std::path::PathBuf::from("/tmp/test_finish_safe.md"));

        // finish_turn with no active turn should not panic
        transcript.finish_turn();

        // Should still work normally after
        transcript.begin_turn(Role::Assistant);
        transcript.stream_delta(BlockType::Text, "Hello");
        transcript.finish_turn();
    }
}