//! Conversation tree module — builds a parent/child tree from `RawMessage`s.
//!
//! Claude Code sessions are not purely linear: the user can edit earlier
//! messages, producing branches. Each message carries a `uuid` and a
//! `parentUuid` that describe the relationship. This module reconstructs
//! the tree so callers can walk threads, detect branches, and select the
//! main thread.

use std::collections::HashMap;

use crate::types::RawMessage;

// ---------------------------------------------------------------------------
// Node & tree types
// ---------------------------------------------------------------------------

/// A single node in the conversation tree.
#[derive(Debug)]
pub struct ConversationNode {
    /// The UUID of this message.
    pub uuid: String,
    /// The raw message stored at this node.
    pub message: RawMessage,
    /// UUIDs of direct children, in insertion order.
    pub children: Vec<String>,
}

/// The full conversation tree for a session.
///
/// A session may have multiple roots when the first message has no
/// `parentUuid`, or when a message refers to a parent that is not present
/// in the slice provided to [`ConversationTree::build`].
#[derive(Debug)]
pub struct ConversationTree {
    /// Root node UUIDs (messages with no parent or with an unknown parent).
    pub roots: Vec<String>,
    /// All nodes indexed by UUID.
    pub nodes: HashMap<String, ConversationNode>,
}

// ---------------------------------------------------------------------------
// UUID / parent-UUID helpers
// ---------------------------------------------------------------------------

/// Extract the `uuid` from any `RawMessage` variant.
///
/// Returns `None` for variants that carry no UUID (e.g. `QueueOperation`).
pub fn message_uuid(msg: &RawMessage) -> Option<&str> {
    match msg {
        RawMessage::User(m) => m.uuid.as_deref(),
        RawMessage::Assistant(m) => m.uuid.as_deref(),
        RawMessage::Progress(m) => m.uuid.as_deref(),
        RawMessage::System(m) => m.uuid.as_deref(),
        RawMessage::QueueOperation(_)
        | RawMessage::FileHistorySnapshot(_)
        | RawMessage::LastPrompt(_) => None,
    }
}

/// Extract the `parent_uuid` from any `RawMessage` variant.
///
/// Returns `None` for variants that carry no parent UUID.
pub fn message_parent_uuid(msg: &RawMessage) -> Option<&str> {
    match msg {
        RawMessage::User(m) => m.parent_uuid.as_deref(),
        RawMessage::Assistant(m) => m.parent_uuid.as_deref(),
        RawMessage::Progress(m) => m.parent_uuid.as_deref(),
        RawMessage::System(m) => m.parent_uuid.as_deref(),
        RawMessage::QueueOperation(_)
        | RawMessage::FileHistorySnapshot(_)
        | RawMessage::LastPrompt(_) => None,
    }
}

// ---------------------------------------------------------------------------
// ConversationTree impl
// ---------------------------------------------------------------------------

impl ConversationTree {
    /// Build a conversation tree from a sequence of messages.
    ///
    /// Messages without a UUID (e.g. `QueueOperation`) are silently skipped.
    /// If a message's `parentUuid` is present but not found in the set,
    /// that message is treated as a root.
    pub fn build(messages: &[RawMessage]) -> Self {
        let mut nodes: HashMap<String, ConversationNode> = HashMap::new();
        let mut roots: Vec<String> = Vec::new();

        // First pass: insert every addressable message as a node.
        for msg in messages {
            if let Some(uuid) = message_uuid(msg) {
                nodes.insert(
                    uuid.to_string(),
                    ConversationNode {
                        uuid: uuid.to_string(),
                        message: msg.clone(),
                        children: Vec::new(),
                    },
                );
            }
        }

        // Second pass: collect (uuid, parent_uuid) pairs so we can wire up
        // parent→child edges without a simultaneous mutable borrow.
        let parent_links: Vec<(String, Option<String>)> = messages
            .iter()
            .filter_map(|msg| {
                let uuid = message_uuid(msg)?.to_string();
                let parent = message_parent_uuid(msg).map(|s| s.to_string());
                Some((uuid, parent))
            })
            .collect();

        for (uuid, parent_uuid) in parent_links {
            match parent_uuid {
                Some(parent_id) if nodes.contains_key(&parent_id) => {
                    // Safe: parent_id != uuid (a message cannot be its own parent).
                    nodes
                        .get_mut(&parent_id)
                        .expect("parent key confirmed above")
                        .children
                        .push(uuid);
                }
                _ => {
                    // No parent, or parent not in the provided slice — treat as root.
                    roots.push(uuid);
                }
            }
        }

        ConversationTree { roots, nodes }
    }

    /// Walk the *main thread*: start from the first root and always follow
    /// the first child at each step.
    ///
    /// In a linear session this is the complete conversation. In a branching
    /// session this is the path taken before any edits.
    pub fn main_thread(&self) -> Vec<&ConversationNode> {
        let mut result = Vec::new();
        if let Some(root_id) = self.roots.first() {
            let mut current = root_id.as_str();
            loop {
                match self.nodes.get(current) {
                    Some(node) => {
                        result.push(node);
                        match node.children.first() {
                            Some(first_child) => current = first_child.as_str(),
                            None => break,
                        }
                    }
                    None => break,
                }
            }
        }
        result
    }

    /// Returns `true` when every node has at most one child (no branches).
    pub fn is_linear(&self) -> bool {
        self.nodes.values().all(|n| n.children.len() <= 1)
    }

    /// Returns all nodes that have more than one child (branch points).
    pub fn branch_points(&self) -> Vec<&ConversationNode> {
        self.nodes
            .values()
            .filter(|n| n.children.len() > 1)
            .collect()
    }
}
