//! Diff widget — wraps [`TreeWidget`] to display a git diff with hunk-level
//! folding and per-line staging marks.
//!
//! Tree structure:
//! ```text
//! Hunk "@@ -10,6 +10,8 @@"   [expandable root node]
//!   ├─ Context  "fn foo() {"  [leaf]
//!   ├─ Added    "+ let x = 1" [leaf]
//!   └─ Removed  "- let y = 2" [leaf]
//! Hunk "@@ -50,3 +52,4 @@"   [expandable root node]
//!   └─ ...
//! ```
//!
//! The `staged` set tracks raw diff indices selected for line-level staging.

use std::collections::HashSet;

use ratatui::{
    buffer::Buffer,
    layout::Rect,
    style::{Color, Modifier, Style},
    text::{Line, Span},
};

use crate::vcs::{DiffKind, DiffLine};
use crate::widgets::tree::{FlatNode, TreeNode, TreeWidget};

// ---------------------------------------------------------------------------
// Data model
// ---------------------------------------------------------------------------

/// Payload for each node in the diff tree.
#[derive(Debug, Clone)]
pub(crate) enum DiffEntry {
    /// A hunk header (`@@` line) — expandable/collapsible parent node.
    Hunk { content: String },
    /// A run of unchanged context lines — collapsible, folded by default.
    ContextBlock { line_count: usize },
    /// A single diff line — leaf node.
    Line { raw_idx: usize, line: DiffLine },
}

// ---------------------------------------------------------------------------
// DiffWidget
// ---------------------------------------------------------------------------

/// Stateful diff widget.  Wraps a `TreeWidget<DiffEntry>` for navigation and
/// fold/expand; owns a `staged` set for line-level staging marks.
#[derive(Debug)]
pub(crate) struct DiffWidget {
    inner: TreeWidget<DiffEntry>,
    /// Raw diff indices of lines the user has toggled for staging.
    staged: HashSet<usize>,
}

impl DiffWidget {
    /// Build the diff tree from a raw `Vec<DiffLine>`.
    ///
    /// Structure:
    /// - Each `HunkHeader` becomes an **expanded** root node.
    /// - Within each hunk, consecutive context lines are grouped into a
    ///   **collapsed** `ContextBlock` inner node.
    /// - Added/removed lines are direct (always-visible) children of the hunk.
    pub(crate) fn new(raw: &[DiffLine]) -> Self {
        Self {
            inner: TreeWidget::new(build_diff_tree(raw)),
            staged: HashSet::new(),
        }
    }

    // -----------------------------------------------------------------------
    // Navigation
    // -----------------------------------------------------------------------

    pub(crate) fn move_up(&mut self) {
        self.inner.move_up();
    }

    pub(crate) fn move_down(&mut self) {
        self.inner.move_down();
    }

    pub(crate) fn page_up(&mut self, n: usize) {
        self.inner.move_up_n(n);
    }

    pub(crate) fn page_down(&mut self, n: usize) {
        self.inner.move_down_n(n);
    }

    // -----------------------------------------------------------------------
    // Actions
    // -----------------------------------------------------------------------

    /// Toggle fold/expand if the cursor is on a `Hunk` node; toggle the
    /// staging mark if the cursor is on a `Line` node.
    pub(crate) fn toggle_stage_or_fold(&mut self) {
        let entry = match self.inner.selected() {
            Some(f) => f.node.data.clone(),
            None => return,
        };
        match entry {
            DiffEntry::Hunk { .. } | DiffEntry::ContextBlock { .. } => {
                self.inner.toggle_selected();
            }
            DiffEntry::Line { raw_idx, .. } => {
                if !self.staged.remove(&raw_idx) {
                    self.staged.insert(raw_idx);
                }
            }
        }
    }

    pub(crate) fn has_staged(&self) -> bool {
        !self.staged.is_empty()
    }

    /// Ordered `(raw_idx, &DiffLine)` pairs for all staged lines.
    /// `raw` must be the original `Vec<DiffLine>` passed to `new`.
    pub(crate) fn staged_lines<'a>(&'a self, raw: &'a [DiffLine]) -> Vec<(usize, &'a DiffLine)> {
        let mut pairs: Vec<(usize, &DiffLine)> = self
            .staged
            .iter()
            .filter_map(|&i| raw.get(i).map(|l| (i, l)))
            .collect();
        pairs.sort_by_key(|(i, _)| *i);
        pairs
    }

    // -----------------------------------------------------------------------
    // Cursor info (used by CommitWindow for mouse handling and editor jump)
    // -----------------------------------------------------------------------

    /// Linear index of the selected node among all currently-visible nodes.
    pub(crate) fn cursor_flat_idx(&self) -> usize {
        self.inner.cursor_flat_idx()
    }

    /// Total number of currently-visible nodes (respecting fold state).
    pub(crate) fn visible_count(&self) -> usize {
        self.inner.flatten().len()
    }

    /// The `line_no` of the currently selected line, if any, for "open in editor".
    pub(crate) fn selected_line_no(&self) -> Option<usize> {
        let flat = self.inner.selected()?;
        match &flat.node.data {
            DiffEntry::Line { line, .. } => line.line_no,
            DiffEntry::Hunk { .. } | DiffEntry::ContextBlock { .. } => None,
        }
    }

    // -----------------------------------------------------------------------
    // Render
    // -----------------------------------------------------------------------

    pub(crate) fn render_into(&self, area: Rect, buf: &mut Buffer, active: bool) {
        let staged = &self.staged;
        self.inner.render(area, buf, |flat, selected| {
            build_line(flat, selected, active, staged)
        });
    }
}

// ---------------------------------------------------------------------------
// Tree builder
// ---------------------------------------------------------------------------

/// Number of context lines to show on each side of a change block.
const CONTEXT_LINES: usize = 3;

/// Build the diff tree with ±`CONTEXT_LINES` visible context around changes.
///
/// Because `vcs` fetches the diff with `context_lines(999_999)`, the raw
/// slice contains the entire file.  We re-select which context lines to
/// display and fold the rest into collapsed `ContextBlock` nodes.
///
/// Two change blocks whose gap is ≤ `2 * CONTEXT_LINES` are automatically
/// merged: the context lines between them are all within range of one side
/// or the other, so no hidden block forms.
fn build_diff_tree(raw: &[DiffLine]) -> Vec<TreeNode<DiffEntry>> {
    // Use the first hunk header for the root display label.
    let hunk_header = raw
        .iter()
        .find(|dl| dl.kind == DiffKind::HunkHeader)
        .map(|dl| dl.content.clone())
        .unwrap_or_default();

    // Flatten to content lines only (skip HunkHeader lines).
    let content: Vec<(usize, &DiffLine)> = raw
        .iter()
        .enumerate()
        .filter(|(_, dl)| dl.kind != DiffKind::HunkHeader)
        .collect();

    let n = content.len();
    if n == 0 {
        return vec![];
    }

    // Positions (in `content`) of every added/removed line.
    let changed: Vec<usize> = content
        .iter()
        .enumerate()
        .filter(|(_, (_, dl))| matches!(dl.kind, DiffKind::Added | DiffKind::Removed))
        .map(|(pos, _)| pos)
        .collect();

    if changed.is_empty() {
        return vec![];
    }

    // Mark visible positions: every change and the CONTEXT_LINES around it.
    let mut visible = vec![false; n];
    for &cp in &changed {
        let lo = cp.saturating_sub(CONTEXT_LINES);
        let hi = (cp + CONTEXT_LINES + 1).min(n);
        for visible_item in visible.iter_mut().take(hi).skip(lo) {
            *visible_item = true;
        }
    }

    // Build children: visible lines are direct leaves; runs of invisible
    // context lines become collapsed ContextBlock inner nodes.
    let mut children: Vec<TreeNode<DiffEntry>> = Vec::new();
    let mut hidden_buf: Vec<TreeNode<DiffEntry>> = Vec::new();

    for (pos, (raw_idx, dl)) in content.iter().enumerate() {
        if visible[pos] {
            if !hidden_buf.is_empty() {
                let count = hidden_buf.len();
                let mut block = TreeNode::inner(
                    DiffEntry::ContextBlock { line_count: count },
                    std::mem::take(&mut hidden_buf),
                );
                block.expanded = false;
                children.push(block);
            }
            children.push(TreeNode::leaf(DiffEntry::Line {
                raw_idx: *raw_idx,
                line: (*dl).clone(),
            }));
        } else {
            hidden_buf.push(TreeNode::leaf(DiffEntry::Line {
                raw_idx: *raw_idx,
                line: (*dl).clone(),
            }));
        }
    }
    if !hidden_buf.is_empty() {
        let count = hidden_buf.len();
        let mut block = TreeNode::inner(
            DiffEntry::ContextBlock { line_count: count },
            hidden_buf,
        );
        block.expanded = false;
        children.push(block);
    }

    let mut root = TreeNode::inner(DiffEntry::Hunk { content: hunk_header }, children);
    root.expanded = true;
    vec![root]
}

// ---------------------------------------------------------------------------
// Line renderer
// ---------------------------------------------------------------------------

fn build_line(
    flat: &FlatNode<'_, DiffEntry>,
    selected: bool,
    active: bool,
    staged: &HashSet<usize>,
) -> Line<'static> {
    match &flat.node.data {
        DiffEntry::Hunk { content } => {
            let icon = if flat.node.expanded { "▾ " } else { "▸ " };
            let text = format!("{}{}", icon, content);
            let style = if selected && active {
                Style::default()
                    .fg(Color::Yellow)
                    .add_modifier(Modifier::BOLD | Modifier::REVERSED)
            } else {
                Style::default()
                    .fg(Color::Yellow)
                    .add_modifier(Modifier::BOLD)
            };
            Line::from(Span::styled(text, style))
        }

        DiffEntry::ContextBlock { line_count } => {
            let icon = if flat.node.expanded { "▾" } else { "▸" };
            let label = if flat.node.expanded {
                format!("{} context", icon)
            } else {
                let s = if *line_count == 1 { "" } else { "s" };
                format!("{} {} unchanged line{}", icon, line_count, s)
            };
            let style = if selected && active {
                Style::default().fg(Color::DarkGray).add_modifier(Modifier::REVERSED)
            } else {
                Style::default().fg(Color::DarkGray)
            };
            Line::from(Span::styled(label, style))
        }

        DiffEntry::Line { raw_idx, line } => {
            let is_staged = staged.contains(raw_idx);
            let sign = match line.kind {
                DiffKind::Added => "+",
                DiffKind::Removed => "-",
                DiffKind::Context => " ",
                DiffKind::HunkHeader => " ", // shouldn't appear as a leaf
            };
            let lineno = line
                .line_no
                .map(|n| format!("{:>4}", n))
                .unwrap_or_else(|| "    ".into());
            let stage_mark = if is_staged { "●" } else { " " };
            let content = format!(
                "{} {} {} {}",
                stage_mark,
                lineno,
                sign,
                truncate(&line.content, 80)
            );

            let base = match line.kind {
                DiffKind::Added => Style::default().fg(Color::Green),
                DiffKind::Removed => Style::default().fg(Color::Red),
                DiffKind::Context | DiffKind::HunkHeader => Style::default().fg(Color::DarkGray),
            };
            let style = match (is_staged, selected && active) {
                (true, true) => {
                    base.add_modifier(Modifier::BOLD | Modifier::UNDERLINED | Modifier::REVERSED)
                }
                (true, false) => base.add_modifier(Modifier::BOLD | Modifier::UNDERLINED),
                (false, true) => base.add_modifier(Modifier::REVERSED),
                (false, false) => base,
            };

            Line::from(Span::styled(content, style))
        }
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn truncate(s: &str, max_chars: usize) -> &str {
    match s.char_indices().nth(max_chars) {
        None => s,
        Some((byte_idx, _)) => &s[..byte_idx],
    }
}


