use anyhow::Result;
use async_trait::async_trait;
use futures::stream::BoxStream;
use serde::{Deserialize, Serialize};

/// Provider-agnostic surface for LLM backends (Ollama, OpenRouter, …).
///
/// Impls translate these canonical shapes at the wire boundary: tool-call
/// arguments stay as `serde_json::Value` in memory and are stringified only
/// when a provider requires it (OpenAI-compatible APIs do), and `images`
/// stays as base64 strings here and is rewritten into content-parts where
/// needed.
///
/// Every method that performs an HTTP call starts an `llm.request` span
/// parented on the caller's ambient `tracing` span, so callers must run
/// inside their feature span (`#[instrument]` / `.instrument(span)`) — and
/// carry it across any `tokio::spawn` — for the LLM span to nest under it.
// First consumer lands in a later PR (OpenRouter impl + hybrid mode routing).
#[allow(dead_code)]
#[async_trait]
pub trait LlmClient: Send + Sync {
    /// Single-shot text generation. Optional system prompt and optional
    /// base64 images (ignored by providers without vision support).
    async fn generate(
        &self,
        prompt: &str,
        system: Option<&str>,
        images: Option<Vec<String>>,
    ) -> Result<String>;

    /// Multi-turn chat with tool definitions. Returns the assistant message
    /// (which may contain tool_calls) plus optional prompt/eval token counts.
    async fn chat_with_tools(
        &self,
        messages: Vec<ChatMessage>,
        tools: Vec<Tool>,
    ) -> Result<(ChatMessage, Option<i32>, Option<i32>)>;

    /// Streaming variant of `chat_with_tools`. The returned stream yields
    /// `TextDelta` items as content is produced, then a single terminal
    /// `Done` carrying the complete assembled message (with tool_calls, if
    /// any) plus token usage counts. Implementations that can't stream may
    /// fall back to calling `chat_with_tools` and emitting the full reply
    /// as one `Done` event.
    async fn chat_with_tools_stream(
        &self,
        messages: Vec<ChatMessage>,
        tools: Vec<Tool>,
    ) -> Result<BoxStream<'static, Result<LlmStreamEvent>>>;

    /// Batch embedding generation. Dimensionality is provider/model specific.
    async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;

    /// One-shot vision description of an image. Used to convert images into
    /// plain text for the hybrid-mode conversation flow.
    async fn describe_image(&self, image_base64: &str) -> Result<String>;

    /// Enumerate available models with their capabilities.
    async fn list_models(&self) -> Result<Vec<ModelCapabilities>>;

    /// Look up capabilities for a single model.
    async fn model_capabilities(&self, model: &str) -> Result<ModelCapabilities>;

    /// Primary model identifier this client was constructed with.
    fn primary_model(&self) -> &str;

    /// Context window (tokens) the server actually serves for the primary
    /// model, when the backend can report one that overrides the request's
    /// `num_ctx`. `None` = unknown or request-controlled (Ollama honours
    /// `num_ctx` per request); callers fall back to the requested value.
    async fn context_window(&self) -> Option<usize> {
        None
    }
}

/// Events emitted by streaming `chat_with_tools_stream`. A stream is a
/// sequence of zero or more `TextDelta` events followed by exactly one
/// `Done`. Callers should treat `Done` as terminal — further items (if any
/// slip through due to upstream misbehavior) are safe to ignore.
#[derive(Debug, Clone)]
pub enum LlmStreamEvent {
    /// Incremental content token(s) from the model. Concatenate in order to
    /// reconstruct the assistant's final text.
    TextDelta(String),
    /// Terminal event with the full assembled message (content + any
    /// tool_calls). `message.content` equals the concatenation of every
    /// preceding `TextDelta.0`.
    Done {
        message: ChatMessage,
        prompt_eval_count: Option<i32>,
        eval_count: Option<i32>,
    },
}

/// Tool definition sent to the model (OpenAI-compatible function schema).
#[derive(Serialize, Clone, Debug)]
pub struct Tool {
    #[serde(rename = "type")]
    pub tool_type: String, // always "function"
    pub function: ToolFunction,
}

#[derive(Serialize, Clone, Debug)]
pub struct ToolFunction {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

impl Tool {
    pub fn function(name: &str, description: &str, parameters: serde_json::Value) -> Self {
        Self {
            tool_type: "function".to_string(),
            function: ToolFunction {
                name: name.to_string(),
                description: description.to_string(),
                parameters,
            },
        }
    }
}

/// A message in the chat conversation history.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ChatMessage {
    pub role: String, // "system" | "user" | "assistant" | "tool"
    /// Empty string (not null) when tool_calls is present — Ollama quirk.
    #[serde(default)]
    pub content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    /// Base64 images — only on user messages to vision-capable models.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub images: Option<Vec<String>>,
}

impl ChatMessage {
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: "system".to_string(),
            content: content.into(),
            tool_calls: None,
            images: None,
        }
    }
    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: "user".to_string(),
            content: content.into(),
            tool_calls: None,
            images: None,
        }
    }
    pub fn tool_result(content: impl Into<String>) -> Self {
        Self {
            role: "tool".to_string(),
            content: content.into(),
            tool_calls: None,
            images: None,
        }
    }
}

/// Tool call returned by the model in an assistant message.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ToolCall {
    pub function: ToolCallFunction,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ToolCallFunction {
    pub name: String,
    /// Canonical shape: native JSON. Providers that use JSON-encoded-string
    /// arguments (OpenAI-compatible) translate at their wire boundary.
    pub arguments: serde_json::Value,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ModelCapabilities {
    pub name: String,
    pub has_vision: bool,
    pub has_tool_calling: bool,
    /// Context window (tokens) the model is served with, when the server
    /// fixes it regardless of the request's `num_ctx` (llama-swap) and it's
    /// known without loading the model. Omitted otherwise.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub context_window: Option<usize>,
}

/// A single node in the conversation tree. Wraps ChatMessage with structural
/// metadata so the tree can be persisted as JSON in `training_messages`.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct StoredChatNode {
    pub id: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parent_id: Option<u64>,
    pub message: ChatMessage,
}

/// Top-level container for tree-based chat history. Stored as JSON in the
/// `training_messages` column. Backward-compatible reader detects whether
/// the stored value is a flat array (old format) or this object (new format).
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ChatHistoryStore {
    pub nodes: Vec<StoredChatNode>,
    pub active_leaf_id: u64,
}

impl ChatHistoryStore {
    /// Create an empty store. Test-only: production stores are always built
    /// from `from_flat_array` or deserialized from `training_messages`.
    #[cfg(test)]
    pub fn empty() -> Self {
        Self {
            nodes: Vec::new(),
            active_leaf_id: 0,
        }
    }

    /// Convert a flat `Vec<ChatMessage>` (old format) into a tree.
    /// Each message becomes a node chained by parent_id.
    pub fn from_flat_array(messages: Vec<ChatMessage>) -> Self {
        let mut nodes = Vec::with_capacity(messages.len());
        let mut prev_id: Option<u64> = None;
        for (i, msg) in messages.into_iter().enumerate() {
            let id = (i + 1) as u64;
            nodes.push(StoredChatNode {
                id,
                parent_id: prev_id,
                message: msg,
            });
            prev_id = Some(id);
        }
        let active_leaf_id = prev_id.unwrap_or(0);
        Self {
            nodes,
            active_leaf_id,
        }
    }

    /// Allocate the next node ID (max existing + 1, or 1 if empty).
    pub fn next_id(&self) -> u64 {
        self.nodes.iter().map(|n| n.id).max().unwrap_or(0) + 1
    }

    /// Append a new node as a child of the given parent. Returns the new node's id.
    pub fn append_node(&mut self, parent_id: Option<u64>, message: ChatMessage) -> u64 {
        let id = self.next_id();
        self.nodes.push(StoredChatNode {
            id,
            parent_id,
            message,
        });
        id
    }

    /// Traverse from root to the given leaf_id, returning the path in order.
    /// Returns None if the leaf_id doesn't exist.
    pub fn path_to_leaf(&self, leaf_id: u64) -> Option<Vec<&StoredChatNode>> {
        // Build a map from id to node for fast lookup.
        let node_map: std::collections::HashMap<u64, usize> = self
            .nodes
            .iter()
            .enumerate()
            .map(|(i, n)| (n.id, i))
            .collect();
        let &_leaf_idx = node_map.get(&leaf_id)?;

        let mut path_indices = Vec::new();
        let mut current_id = leaf_id;
        while let Some(&idx) = node_map.get(&current_id) {
            path_indices.push(idx);
            match self.nodes[idx].parent_id {
                Some(pid) => current_id = pid,
                None => break,
            }
        }
        path_indices.reverse();
        Some(path_indices.into_iter().map(|i| &self.nodes[i]).collect())
    }

    /// All direct children of a given node, ordered by id (chronological).
    pub fn children_of(&self, node_id: u64) -> Vec<&StoredChatNode> {
        self.nodes
            .iter()
            .filter(|n| n.parent_id == Some(node_id))
            .collect()
    }

    /// For a given node, return (position, total) where:
    /// - total = number of children of the parent of this node (i.e., siblings + self)
    /// - position = 1-based rank of this node among siblings, ordered by id
    ///
    /// Returns None if the node has no parent (root) or the parent has 0 or 1 child.
    pub fn fork_at_node(&self, node_id: u64) -> Option<(usize, usize)> {
        let node = self.nodes.iter().find(|n| n.id == node_id)?;
        let parent_id = node.parent_id?;
        let children = self.children_of(parent_id);
        let total = children.len();
        if total <= 1 {
            return None;
        }
        let position = children.iter().position(|c| c.id == node_id)? + 1;
        Some((position, total))
    }

    /// All leaf node IDs within the subtree rooted at `node_id` (the node
    /// itself when it has no children). Empty when the node doesn't exist.
    pub fn descendant_leaves(&self, node_id: u64) -> Vec<u64> {
        if !self.nodes.iter().any(|n| n.id == node_id) {
            return Vec::new();
        }
        let mut result = Vec::new();
        let mut stack = vec![node_id];
        while let Some(id) = stack.pop() {
            let children = self.children_of(id);
            if children.is_empty() {
                result.push(id);
            } else {
                stack.extend(children.iter().map(|c| c.id));
            }
        }
        result
    }

    /// Branch options at a divergence point: one entry per direct child of
    /// `fork_node_id`, ordered by id (chronological — matches the stable
    /// position rank shown in "X/Y" fork chips).
    ///
    /// Each entry's `id` is the leaf to switch to for that branch: the
    /// `preferred_leaf` when it lives in that child's subtree (so the
    /// caller's current view maps onto its own option and can be filtered
    /// out client-side), otherwise the most recent (max-id) descendant leaf.
    ///
    /// Snippets preview the first visible message that *differs* between
    /// the branches, not blindly the first child: Rewind & Regenerate
    /// resends the identical question, so the first child is often shared
    /// across siblings and would make every option read the same.
    pub fn branch_options_at(&self, fork_node_id: u64, preferred_leaf: u64) -> Vec<BranchLeafInfo> {
        struct Candidate {
            rep: u64,
            message_count: usize,
            /// User/assistant contents after the divergence, tool
            /// scaffolding (tool results, empty dispatch turns) skipped.
            visible: Vec<String>,
        }
        let candidates: Vec<Candidate> = self
            .children_of(fork_node_id)
            .iter()
            .map(|child| {
                let leaves = self.descendant_leaves(child.id);
                let rep = if leaves.contains(&preferred_leaf) {
                    preferred_leaf
                } else {
                    leaves.iter().copied().max().unwrap_or(child.id)
                };
                let path = self.path_to_leaf(rep).unwrap_or_default();
                let message_count = path.len();
                let visible: Vec<String> = path
                    .iter()
                    .skip_while(|n| n.id != child.id)
                    .filter(|n| {
                        (n.message.role == "user" || n.message.role == "assistant")
                            && !n.message.content.trim().is_empty()
                    })
                    .map(|n| n.message.content.clone())
                    .collect();
                Candidate {
                    rep,
                    message_count,
                    visible,
                }
            })
            .collect();

        // Walk forward past positions where every branch has the same
        // content; stop at the first divergence (or when any branch runs
        // out — a length difference is itself distinguishing).
        let mut snippet_idx = 0usize;
        loop {
            let contents: Vec<Option<&String>> = candidates
                .iter()
                .map(|c| c.visible.get(snippet_idx))
                .collect();
            let all_present_and_equal =
                contents.iter().all(|c| c.is_some()) && contents.windows(2).all(|w| w[0] == w[1]);
            if !all_present_and_equal {
                break;
            }
            snippet_idx += 1;
        }

        candidates
            .into_iter()
            .enumerate()
            .map(|(i, c)| {
                // A branch exhausted before the divergence index falls back
                // to its own last message so it still shows *something*.
                let snippet = c
                    .visible
                    .get(snippet_idx)
                    .or_else(|| c.visible.last())
                    .map(|s| s.chars().take(100).collect::<String>())
                    .unwrap_or_default();
                BranchLeafInfo {
                    id: c.rep,
                    snippet,
                    message_count: c.message_count,
                    position: Some(i + 1),
                }
            })
            .collect()
    }

    /// All leaf node IDs (nodes with no children).
    pub fn leaves(&self) -> Vec<u64> {
        let child_ids: std::collections::HashSet<u64> =
            self.nodes.iter().filter_map(|n| n.parent_id).collect();
        self.nodes
            .iter()
            .filter(|n| !child_ids.contains(&n.id))
            .map(|n| n.id)
            .collect()
    }

    /// The parent of the given node, if any. Used to re-anchor a rewind
    /// that discards every rendered message onto the seed node above them.
    pub fn parent_of(&self, node_id: u64) -> Option<&StoredChatNode> {
        let node = self.nodes.iter().find(|n| n.id == node_id)?;
        let parent_id = node.parent_id?;
        self.nodes.iter().find(|n| n.id == parent_id)
    }

    /// Return all leaf nodes with metadata: id, snippet (first 100 chars of
    /// the first user-visible message after the divergence point), and
    /// message_count (number of nodes in the path from root to leaf).
    pub fn leaves_with_info(&self) -> Vec<BranchLeafInfo> {
        let leaf_ids = self.leaves();
        leaf_ids
            .into_iter()
            .map(|leaf_id| {
                let path = self.path_to_leaf(leaf_id).unwrap_or_default();
                let message_count = path.len();
                // Snippet: first 100 chars of the first user/assistant message
                // (skip system messages).
                let snippet = path
                    .iter()
                    .find(|n| n.message.role != "system")
                    .map(|n| n.message.content.chars().take(100).collect::<String>())
                    .unwrap_or_default();
                BranchLeafInfo {
                    id: leaf_id,
                    snippet,
                    message_count,
                    position: None,
                }
            })
            .collect()
    }
}

/// Metadata about a branch leaf, returned by the branches endpoint.
#[derive(Serialize, Clone, Debug)]
pub struct BranchLeafInfo {
    pub id: u64,
    #[serde(skip_serializing_if = "String::is_empty")]
    pub snippet: String,
    pub message_count: usize,
    /// 1-based sibling rank at the divergence point. Only present for
    /// scoped queries (`branch_options_at`); tree-wide leaf listings have
    /// no single divergence to rank against.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub position: Option<usize>,
}

/// Strip a leading `<think>…</think>` reasoning block from model output.
///
/// Thinking models sometimes emit chain-of-thought inside think tags before
/// the real answer. Everything after the first `</think>` is the answer;
/// when no tag is present — or the text after it is empty — the trimmed
/// input is returned unchanged. Mirrors the behavior Ollama's
/// `extract_final_answer` has applied to single-shot generation; shared here
/// so the tool-calling final-content paths (agentic generation + chat) can
/// apply the identical cleanup before parsing / persisting.
pub fn strip_think_blocks(response: &str) -> String {
    let response = response.trim();

    if let Some(pos) = response.find("</think>") {
        let answer = response[pos + "</think>".len()..].trim();
        if !answer.is_empty() {
            return answer.to_string();
        }
    }

    response.to_string()
}

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

    #[test]
    fn strip_think_blocks_removes_leading_think_block() {
        let raw = "<think>\nLet me reason about this.\n</think>\n\nTitle: A Day Out\n\nThe body.";
        assert_eq!(strip_think_blocks(raw), "Title: A Day Out\n\nThe body.");
    }

    #[test]
    fn strip_think_blocks_passes_through_plain_content() {
        assert_eq!(strip_think_blocks("  just an answer  "), "just an answer");
    }

    #[test]
    fn strip_think_blocks_keeps_content_when_answer_after_tag_is_empty() {
        // A think block with nothing after it: better to return the trimmed
        // original than an empty string (matches Ollama's fallback).
        let raw = "<think>only thoughts</think>";
        assert_eq!(strip_think_blocks(raw), raw);
    }

    #[test]
    fn strip_think_blocks_handles_unclosed_tag() {
        let raw = "<think>thinking forever";
        assert_eq!(strip_think_blocks(raw), raw);
    }

    // ── ChatHistoryStore tree traversal tests ──────────────────────────

    fn mk_msg(role: &str, content: &str) -> ChatMessage {
        match role {
            "system" => ChatMessage::system(content),
            "user" => ChatMessage::user(content),
            _ => ChatMessage {
                role: role.to_string(),
                content: content.to_string(),
                tool_calls: None,
                images: None,
            },
        }
    }

    #[test]
    fn store_from_flat_array_chains_parent_ids() {
        let msgs = vec![
            mk_msg("system", "sys"),
            mk_msg("user", "hello"),
            mk_msg("assistant", "hi there"),
        ];
        let store = ChatHistoryStore::from_flat_array(msgs);
        assert_eq!(store.nodes.len(), 3);
        assert_eq!(store.nodes[0].id, 1);
        assert_eq!(store.nodes[0].parent_id, None);
        assert_eq!(store.nodes[1].parent_id, Some(1));
        assert_eq!(store.nodes[2].parent_id, Some(2));
        assert_eq!(store.active_leaf_id, 3);
    }

    #[test]
    fn store_path_to_leaf_returns_ordered_path() {
        let mut store = ChatHistoryStore::empty();
        let n1 = store.append_node(None, mk_msg("user", "u1"));
        let n2 = store.append_node(Some(n1), mk_msg("assistant", "a1"));
        let n3 = store.append_node(Some(n2), mk_msg("user", "u2"));
        let path = store.path_to_leaf(n3).unwrap();
        assert_eq!(path.len(), 3);
        assert_eq!(path[0].id, n1);
        assert_eq!(path[1].id, n2);
        assert_eq!(path[2].id, n3);
    }

    #[test]
    fn store_path_to_leaf_returns_none_for_missing_id() {
        let store = ChatHistoryStore::empty();
        assert!(store.path_to_leaf(999).is_none());
    }

    #[test]
    fn store_children_of_returns_direct_children() {
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "root"));
        let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
        let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
        let _d1 = store.append_node(Some(c1), mk_msg("assistant", "d1"));
        let children = store.children_of(root);
        assert_eq!(children.len(), 2);
        assert_eq!(children[0].id, c1);
        assert_eq!(children[1].id, c2);
    }

    #[test]
    fn store_fork_at_node_detects_fork() {
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "root"));
        let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
        let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
        let c3 = store.append_node(Some(root), mk_msg("user", "c3"));
        // c1 is position 1 of 3 siblings
        assert_eq!(store.fork_at_node(c1), Some((1, 3)));
        assert_eq!(store.fork_at_node(c2), Some((2, 3)));
        assert_eq!(store.fork_at_node(c3), Some((3, 3)));
        // root has no parent, so no fork
        assert!(store.fork_at_node(root).is_none());
    }

    #[test]
    fn store_fork_at_node_returns_none_for_single_child() {
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "root"));
        let _only = store.append_node(Some(root), mk_msg("assistant", "only"));
        assert!(store.fork_at_node(2).is_none());
    }

    #[test]
    fn store_leaves_returns_leaf_ids() {
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "root"));
        let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
        let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
        let _d1 = store.append_node(Some(c1), mk_msg("assistant", "d1"));
        let leaves = store.leaves();
        assert_eq!(leaves.len(), 2);
        assert!(leaves.contains(&c2));
        assert!(leaves.contains(&_d1));
        assert!(!leaves.contains(&root));
        assert!(!leaves.contains(&c1));
    }

    #[test]
    fn store_parent_of_returns_parent() {
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "root"));
        let child = store.append_node(Some(root), mk_msg("assistant", "child"));
        let parent = store.parent_of(child).unwrap();
        assert_eq!(parent.id, root);
        assert!(store.parent_of(root).is_none());
    }

    #[test]
    fn store_next_id_increments() {
        let mut store = ChatHistoryStore::empty();
        assert_eq!(store.next_id(), 1);
        store.append_node(None, mk_msg("user", "a"));
        assert_eq!(store.next_id(), 2);
        store.append_node(Some(1), mk_msg("assistant", "b"));
        assert_eq!(store.next_id(), 3);
    }

    #[test]
    fn store_descendant_leaves_walks_subtree() {
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "root"));
        let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
        let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
        let d1 = store.append_node(Some(c1), mk_msg("assistant", "d1"));
        let d2 = store.append_node(Some(c1), mk_msg("assistant", "d2"));
        // c1's subtree has two leaves; c2 is its own leaf; root sees all.
        let mut c1_leaves = store.descendant_leaves(c1);
        c1_leaves.sort_unstable();
        assert_eq!(c1_leaves, vec![d1, d2]);
        assert_eq!(store.descendant_leaves(c2), vec![c2]);
        let mut all = store.descendant_leaves(root);
        all.sort_unstable();
        assert_eq!(all, vec![c2, d1, d2]);
        assert!(store.descendant_leaves(999).is_empty());
    }

    #[test]
    fn branch_options_ranks_children_and_prefers_viewing_leaf() {
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "root"));
        // First branch: two leaves (older + newer).
        let c1 = store.append_node(Some(root), mk_msg("assistant", "reply one"));
        let old_leaf = store.append_node(Some(c1), mk_msg("user", "old follow-up"));
        let new_leaf = store.append_node(Some(c1), mk_msg("user", "new follow-up"));
        // Second branch: single leaf.
        let c2 = store.append_node(Some(root), mk_msg("assistant", "reply two"));

        // Viewing the *older* leaf of branch 1: its option must anchor to
        // that exact leaf, not the max-id one, so the client's
        // "hide the branch I'm on" filter matches.
        let opts = store.branch_options_at(root, old_leaf);
        assert_eq!(opts.len(), 2);
        assert_eq!(opts[0].id, old_leaf);
        assert_eq!(opts[0].position, Some(1));
        assert_eq!(opts[0].snippet, "reply one");
        assert_eq!(opts[1].id, c2);
        assert_eq!(opts[1].position, Some(2));

        // Viewing branch 2: branch 1 falls back to its most recent leaf.
        let opts = store.branch_options_at(root, c2);
        assert_eq!(opts[0].id, new_leaf);
        assert_eq!(opts[1].id, c2);
    }

    #[test]
    fn branch_options_snippet_skips_shared_resent_question() {
        // Rewind & Regenerate resends the identical question, so both
        // branches' first child matches — the preview must advance to the
        // first *differing* message (the assistant replies) or every
        // option would read the same.
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("assistant", "prior reply"));
        let q1 = store.append_node(Some(root), mk_msg("user", "same question"));
        let _a1 = store.append_node(Some(q1), mk_msg("assistant", "first answer"));
        let q2 = store.append_node(Some(root), mk_msg("user", "same question"));
        let a2 = store.append_node(Some(q2), mk_msg("assistant", "second answer"));

        let opts = store.branch_options_at(root, a2);
        assert_eq!(opts.len(), 2);
        assert_eq!(opts[0].snippet, "first answer");
        assert_eq!(opts[1].snippet, "second answer");
    }

    #[test]
    fn branch_options_snippet_falls_back_when_branch_exhausted() {
        // One branch is just the shared question (no reply yet); the other
        // continues past it. The shorter branch falls back to its own last
        // message rather than showing an empty preview.
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("assistant", "prior reply"));
        let _q1 = store.append_node(Some(root), mk_msg("user", "same question"));
        let q2 = store.append_node(Some(root), mk_msg("user", "same question"));
        let a2 = store.append_node(Some(q2), mk_msg("assistant", "second answer"));

        let opts = store.branch_options_at(root, a2);
        assert_eq!(opts[0].snippet, "same question");
        assert_eq!(opts[1].snippet, "second answer");
    }

    #[test]
    fn branch_options_snippet_skips_tool_scaffolding() {
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "root"));
        // Branch whose diverging child is an empty tool-dispatch turn.
        let dispatch = store.append_node(
            Some(root),
            ChatMessage {
                role: "assistant".to_string(),
                content: String::new(),
                tool_calls: None,
                images: None,
            },
        );
        let tool = store.append_node(Some(dispatch), mk_msg("tool", "raw tool output"));
        let reply = store.append_node(Some(tool), mk_msg("assistant", "final answer"));
        let _other = store.append_node(Some(root), mk_msg("assistant", "direct reply"));

        let opts = store.branch_options_at(root, reply);
        assert_eq!(opts[0].snippet, "final answer");
        assert_eq!(opts[1].snippet, "direct reply");
    }

    #[test]
    fn store_append_node_sets_parent() {
        let mut store = ChatHistoryStore::empty();
        let parent = store.append_node(None, mk_msg("user", "parent"));
        let child = store.append_node(Some(parent), mk_msg("assistant", "child"));
        assert_eq!(store.nodes.len(), 2);
        assert_eq!(store.nodes[1].parent_id, Some(parent));
        assert_eq!(store.nodes[1].id, child);
    }

    #[test]
    fn branch_options_handles_nested_forks() {
        // Regression: verify sub-fork display when there are multiple forks.
        // Structure:
        //   root (user "q1")
        //   ├── a1 (assistant "answer 1")
        //   └── a2 (assistant "answer 2")
        //       ├── q2 (user "q2")
        //       │   ├── r1 (assistant "reply 1")
        //       │   └── r2 (assistant "reply 2")  <- active leaf
        //
        // branch_options_at(root) should return [a1, a2] (2 branches)
        // branch_options_at(q2) should return [r1, r2] (2 sub-branches)
        let mut store = ChatHistoryStore::empty();
        let root = store.append_node(None, mk_msg("user", "q1"));
        let a1 = store.append_node(Some(root), mk_msg("assistant", "answer 1"));
        let a2 = store.append_node(Some(root), mk_msg("assistant", "answer 2"));
        let q2 = store.append_node(Some(a2), mk_msg("user", "q2"));
        let r1 = store.append_node(Some(q2), mk_msg("assistant", "reply 1"));
        let r2 = store.append_node(Some(q2), mk_msg("assistant", "reply 2"));
        store.active_leaf_id = r2;

        // Root fork: should show [a1, a2's subtree]
        let opts = store.branch_options_at(root, r2);
        assert_eq!(opts.len(), 2);
        assert_eq!(opts[0].id, a1);
        assert_eq!(opts[0].position, Some(1));
        assert_eq!(opts[0].snippet, "answer 1");
        // a2's subtree includes r2, so the rep should be r2 (preferred leaf)
        assert_eq!(opts[1].id, r2);
        assert_eq!(opts[1].position, Some(2));
        assert_eq!(opts[1].snippet, "answer 2");

        // Sub-fork at q2: should show [r1, r2]
        let opts = store.branch_options_at(q2, r2);
        assert_eq!(opts.len(), 2);
        assert_eq!(opts[0].id, r1);
        assert_eq!(opts[0].position, Some(1));
        assert_eq!(opts[0].snippet, "reply 1");
        assert_eq!(opts[1].id, r2);
        assert_eq!(opts[1].position, Some(2));
        assert_eq!(opts[1].snippet, "reply 2");
    }
}