travelagent 1.11.1

use std::collections::HashSet;
use std::path::Path;

use travelagent_core::error::{Result, TrvError};
use travelagent_core::model::{DiffLine, LineOrigin, LineSide};
use travelagent_core::vcs::git::calculate_gap;

use super::{
    AnnotatedLine, App, DiffViewMode, ExpandDirection, FileTreeItem, GAP_EXPAND_BATCH,
    GapCursorHit, GapId,
};

impl App {
    pub fn expand_all_dirs(&mut self) {
        self.ui_layout.expanded_dirs.clear();
        for file in &self.diff_files {
            let path = file.display_path_lossy();
            let mut current = path.parent();
            while let Some(parent) = current {
                if parent != Path::new("") {
                    self.ui_layout
                        .expanded_dirs
                        .insert(parent.to_string_lossy().to_string());
                }
                current = parent.parent();
            }
        }
        self.ensure_valid_tree_selection();
    }

    pub fn collapse_all_dirs(&mut self) {
        self.ui_layout.expanded_dirs.clear();
        self.ensure_valid_tree_selection();
    }

    pub fn toggle_directory(&mut self, dir_path: &str) {
        if self.ui_layout.expanded_dirs.contains(dir_path) {
            self.ui_layout.expanded_dirs.remove(dir_path);
            // After collapsing, keep the cursor on the directory item itself
            // rather than letting ensure_valid_tree_selection move it elsewhere.
            let visible_items = self.build_visible_items();
            if let Some(dir_idx) = visible_items.iter().position(
                |item| matches!(item, FileTreeItem::Directory { path, .. } if path == dir_path),
            ) {
                self.file_list_state.select(dir_idx);
            } else {
                self.ensure_valid_tree_selection();
            }
        } else {
            self.ui_layout.expanded_dirs.insert(dir_path.to_string());
        }
    }

    /// Get the line boundaries (start_line, end_line) of a gap.
    fn gap_boundaries(&self, gap_id: &GapId) -> Option<(u32, u32)> {
        let file = self.diff_files.get(gap_id.file_idx)?;
        let hunk = file.hunks.get(gap_id.hunk_idx)?;
        let prev_hunk = if gap_id.hunk_idx > 0 {
            file.hunks.get(gap_id.hunk_idx - 1)
        } else {
            None
        };
        let (start, end) = match prev_hunk {
            None => (1, hunk.new_start.saturating_sub(1)),
            Some(prev) => (
                prev.new_start + prev.new_count,
                hunk.new_start.saturating_sub(1),
            ),
        };
        if start > end {
            None
        } else {
            Some((start, end))
        }
    }

    /// Look up an expanded context line by sequential index across top + bottom.
    pub(super) fn get_expanded_line(&self, gap_id: &GapId, idx: usize) -> Option<&DiffLine> {
        let top = self.gaps.expanded_top.get(gap_id);
        let top_len = top.map_or(0, std::vec::Vec::len);
        if idx < top_len {
            top?.get(idx)
        } else {
            self.gaps.expanded_bottom.get(gap_id)?.get(idx - top_len)
        }
    }

    /// Expand a gap in the given direction.
    /// If `limit` is Some(n), expand up to n lines. If None, expand all remaining.
    pub fn expand_gap(
        &mut self,
        gap_id: GapId,
        direction: ExpandDirection,
        limit: Option<usize>,
    ) -> Result<()> {
        let (gap_start, gap_end) = self
            .gap_boundaries(&gap_id)
            .ok_or_else(|| TrvError::CorruptedSession(format!("Invalid gap: {gap_id:?}")))?;

        let file_path = self.diff_files[gap_id.file_idx]
            .display_path_lossy()
            .clone();
        let file_status = self.diff_files[gap_id.file_idx].status;

        let top_len = self
            .gaps
            .expanded_top
            .get(&gap_id)
            .map_or(0, std::vec::Vec::len) as u32;
        let bot_len = self
            .gaps
            .expanded_bottom
            .get(&gap_id)
            .map_or(0, std::vec::Vec::len) as u32;

        // The unexpanded region runs from (gap_start + top_len) to (gap_end - bot_len)
        let inner_start = gap_start + top_len;
        let inner_end = gap_end.saturating_sub(bot_len);

        if inner_start > inner_end {
            return Ok(()); // Fully expanded
        }

        match direction {
            ExpandDirection::Down => {
                let n = limit.unwrap_or(usize::MAX) as u32;
                if n == 0 {
                    return Ok(());
                }
                let fetch_end = inner_start
                    .saturating_add(n.saturating_sub(1))
                    .min(inner_end);
                let new_lines = self.vcs.fetch_context_lines(
                    &file_path,
                    file_status,
                    inner_start,
                    fetch_end,
                )?;
                self.gaps
                    .expanded_top
                    .entry(gap_id.clone())
                    .or_default()
                    .extend(new_lines);
            }
            ExpandDirection::Up => {
                let n = limit.unwrap_or(usize::MAX) as u32;
                if n == 0 {
                    return Ok(());
                }
                let fetch_start = inner_end
                    .saturating_sub(n.saturating_sub(1))
                    .max(inner_start);
                let new_lines = self.vcs.fetch_context_lines(
                    &file_path,
                    file_status,
                    fetch_start,
                    inner_end,
                )?;
                // Prepend: new lines go before existing bottom lines
                let existing = self
                    .gaps
                    .expanded_bottom
                    .remove(&gap_id)
                    .unwrap_or_default();
                let mut combined = new_lines;
                combined.extend(existing);
                self.gaps.expanded_bottom.insert(gap_id.clone(), combined);
            }
            ExpandDirection::Both => {
                // Fetch everything remaining
                let new_lines = self.vcs.fetch_context_lines(
                    &file_path,
                    file_status,
                    inner_start,
                    inner_end,
                )?;
                self.gaps
                    .expanded_top
                    .entry(gap_id.clone())
                    .or_default()
                    .extend(new_lines);
            }
        }

        self.rebuild_annotations();
        Ok(())
    }

    /// Collapse an expanded gap
    pub fn collapse_gap(&mut self, gap_id: GapId) {
        self.gaps.expanded_top.remove(&gap_id);
        self.gaps.expanded_bottom.remove(&gap_id);
        self.rebuild_annotations();
    }

    /// Clear all expanded gaps (called when reloading diffs)
    pub fn clear_expanded_gaps(&mut self) {
        self.gaps.expanded_top.clear();
        self.gaps.expanded_bottom.clear();
    }

    /// Rebuild the line annotations cache. Call this when:
    /// - Diff files change (load/reload)
    /// - Expansion state changes (expand/collapse gap)
    /// - Comments are added/removed
    /// - Diff view mode changes
    pub fn rebuild_annotations(&mut self) {
        self.line_annotations.clear();

        self.line_annotations
            .push(AnnotatedLine::ReviewCommentsHeader);
        for (comment_idx, comment) in self.engine.session().review_comments.iter().enumerate() {
            let comment_lines = Self::comment_display_lines(comment);
            for _ in 0..comment_lines {
                self.line_annotations
                    .push(AnnotatedLine::ReviewComment { comment_idx });
            }
        }

        // L3: session-level orphaned comments — orphans living on
        // `FileReview`s whose path is no longer in `diff_files`. These
        // appear above the per-file sections so the user can see and
        // re-anchor them before diving into the fresh diff.
        //
        // TODO(live-mode): collapse after N > 10 orphans — see
        // docs/design/live-review-mode.md "Risks + open questions".
        let diff_paths: std::collections::HashSet<_> = self
            .diff_files
            .iter()
            .map(|f| f.display_path_lossy().clone())
            .collect();

        // Single-pass collection of gone files + their per-orphan display
        // heights. `session.files`'s HashMap order is deterministic within
        // one rebuild but not across runs — sort by path for display
        // stability. We capture `comment_display_lines` up-front so the
        // header-building loop below does zero extra HashMap lookups.
        let mut gone: Vec<(std::path::PathBuf, Vec<usize>)> = self
            .engine
            .session()
            .files
            .iter()
            .filter(|(p, _)| !diff_paths.contains(*p))
            .map(|(p, r)| {
                (
                    p.clone(),
                    r.orphaned_comments
                        .iter()
                        .map(Self::comment_display_lines)
                        .collect::<Vec<_>>(),
                )
            })
            .collect();
        gone.sort_by(|a, b| a.0.cmp(&b.0));
        let session_orphan_total: usize = gone.iter().map(|(_, lines)| lines.len()).sum();
        if session_orphan_total > 0 {
            self.line_annotations
                .push(AnnotatedLine::OrphanedCommentsHeader {
                    file_idx: None,
                    count: session_orphan_total,
                });
            for (path, per_orphan_lines) in gone {
                for (orphan_idx, lines) in per_orphan_lines.into_iter().enumerate() {
                    // +1 for the per-orphan header line showing "was line N"
                    for _ in 0..(lines + 1) {
                        self.line_annotations.push(AnnotatedLine::OrphanedComment {
                            file_idx: None,
                            orphan_idx,
                            file_path: path.clone(),
                        });
                    }
                }
            }
            self.line_annotations.push(AnnotatedLine::Spacing);
        }

        for (file_idx, file) in self.diff_files.iter().enumerate() {
            let path = file.display_path_lossy();

            // L3: per-file orphaned comments — before the file header, so
            // the user sees the orphans near the rest of the file's state.
            if let Some(review) = self.engine.session().files.get(path.as_path())
                && !review.orphaned_comments.is_empty()
            {
                self.line_annotations
                    .push(AnnotatedLine::OrphanedCommentsHeader {
                        file_idx: Some(file_idx),
                        count: review.orphaned_comments.len(),
                    });
                for (orphan_idx, comment) in review.orphaned_comments.iter().enumerate() {
                    let lines = Self::comment_display_lines(comment);
                    // +1 for the per-orphan "was line N" header line.
                    for _ in 0..(lines + 1) {
                        self.line_annotations.push(AnnotatedLine::OrphanedComment {
                            file_idx: Some(file_idx),
                            orphan_idx,
                            file_path: path.clone(),
                        });
                    }
                }
                self.line_annotations.push(AnnotatedLine::Spacing);
            }

            // File header
            self.line_annotations
                .push(AnnotatedLine::FileHeader { file_idx });

            // If reviewed, skip all content for this file
            if self.engine.session().is_file_reviewed(path) {
                continue;
            }

            // If collapsed (either auto-collapsed because the file is
            // lockfile-like / oversized, or explicitly collapsed via `z`),
            // emit a single placeholder line and skip the body.
            if self.is_file_collapsed(file_idx) {
                self.line_annotations
                    .push(AnnotatedLine::CollapsedFile { file_idx });
                self.line_annotations.push(AnnotatedLine::Spacing);
                continue;
            }

            // File comments
            if let Some(review) = self.engine.session().files.get(path) {
                for (comment_idx, comment) in review.file_comments.iter().enumerate() {
                    let comment_lines = Self::comment_display_lines(comment);
                    for _ in 0..comment_lines {
                        self.line_annotations.push(AnnotatedLine::FileComment {
                            file_idx,
                            comment_idx,
                        });
                    }
                }
            }

            if file.is_binary || file.hunks.is_empty() {
                self.line_annotations
                    .push(AnnotatedLine::BinaryOrEmpty { file_idx });
            } else {
                // Get line comments for this file
                let line_comments = self.engine.line_comments_for(path);

                for (hunk_idx, hunk) in file.hunks.iter().enumerate() {
                    // Calculate gap before this hunk
                    let prev_hunk = if hunk_idx > 0 {
                        file.hunks.get(hunk_idx - 1)
                    } else {
                        None
                    };
                    let gap = calculate_gap(
                        prev_hunk.map(|h| (&h.new_start, &h.new_count)),
                        hunk.new_start,
                    );

                    let gap_id = GapId { file_idx, hunk_idx };

                    if gap > 0 {
                        let top_len = self
                            .gaps
                            .expanded_top
                            .get(&gap_id)
                            .map_or(0, std::vec::Vec::len);
                        let bot_len = self
                            .gaps
                            .expanded_bottom
                            .get(&gap_id)
                            .map_or(0, std::vec::Vec::len);
                        let remaining = (gap as usize).saturating_sub(top_len + bot_len);
                        let is_top_of_file = hunk_idx == 0;

                        // Sequential line_idx counter across top + bottom
                        let mut ctx_idx = 0;

                        // --- Top expanded lines (↓ direction) ---
                        for _ in 0..top_len {
                            self.line_annotations.push(AnnotatedLine::ExpandedContext {
                                gap_id: gap_id.clone(),
                                line_idx: ctx_idx,
                            });
                            ctx_idx += 1;
                        }

                        // --- Expanders / hidden lines ---
                        if remaining > 0 {
                            if is_top_of_file {
                                // Top-of-file: HiddenLines (if > batch) + ↑
                                if remaining > GAP_EXPAND_BATCH {
                                    self.line_annotations.push(AnnotatedLine::HiddenLines {
                                        gap_id: gap_id.clone(),
                                        count: remaining,
                                    });
                                }
                                self.line_annotations.push(AnnotatedLine::Expander {
                                    gap_id: gap_id.clone(),
                                    direction: ExpandDirection::Up,
                                });
                            } else if remaining >= GAP_EXPAND_BATCH {
                                // Between-hunk, large: ↓ + HiddenLines + ↑
                                self.line_annotations.push(AnnotatedLine::Expander {
                                    gap_id: gap_id.clone(),
                                    direction: ExpandDirection::Down,
                                });
                                self.line_annotations.push(AnnotatedLine::HiddenLines {
                                    gap_id: gap_id.clone(),
                                    count: remaining,
                                });
                                self.line_annotations.push(AnnotatedLine::Expander {
                                    gap_id: gap_id.clone(),
                                    direction: ExpandDirection::Up,
                                });
                            } else {
                                // Between-hunk, small: merged ↕
                                self.line_annotations.push(AnnotatedLine::Expander {
                                    gap_id: gap_id.clone(),
                                    direction: ExpandDirection::Both,
                                });
                            }
                        }

                        // --- Bottom expanded lines (↑ direction) ---
                        for _ in 0..bot_len {
                            self.line_annotations.push(AnnotatedLine::ExpandedContext {
                                gap_id: gap_id.clone(),
                                line_idx: ctx_idx,
                            });
                            ctx_idx += 1;
                        }
                    }

                    // Hunk header
                    self.line_annotations
                        .push(AnnotatedLine::HunkHeader { file_idx, hunk_idx });

                    // Diff lines - handle differently based on view mode
                    match self.nav.diff_view_mode {
                        DiffViewMode::Unified => {
                            Self::build_unified_diff_annotations(
                                &mut self.line_annotations,
                                file_idx,
                                hunk_idx,
                                &hunk.lines,
                                &line_comments,
                            );
                        }
                        DiffViewMode::SideBySide => {
                            Self::build_side_by_side_annotations(
                                &mut self.line_annotations,
                                file_idx,
                                hunk_idx,
                                &hunk.lines,
                                &line_comments,
                            );
                        }
                    }
                }
            }

            // Spacing line
            self.line_annotations.push(AnnotatedLine::Spacing);
        }
    }

    fn push_comments(
        annotations: &mut Vec<AnnotatedLine>,
        file_idx: usize,
        line_no: Option<u32>,
        line_comments: &std::collections::HashMap<u32, Vec<travelagent_core::model::Comment>>,
        side: LineSide,
    ) {
        let Some(ln) = line_no else {
            return;
        };

        let Some(comments) = line_comments.get(&ln) else {
            return;
        };

        for (idx, comment) in comments.iter().enumerate() {
            let matches_side =
                comment.side == Some(side) || (side == LineSide::New && comment.side.is_none());

            if !matches_side {
                continue;
            }

            let comment_lines = Self::comment_display_lines(comment);
            for _ in 0..comment_lines {
                annotations.push(AnnotatedLine::LineComment {
                    file_idx,
                    line: ln,
                    comment_idx: idx,
                    side,
                });
            }
        }
    }

    /// Build annotations for unified diff mode (one annotation per diff line)
    fn build_unified_diff_annotations(
        annotations: &mut Vec<AnnotatedLine>,
        file_idx: usize,
        hunk_idx: usize,
        lines: &[travelagent_core::model::DiffLine],
        line_comments: &std::collections::HashMap<u32, Vec<travelagent_core::model::Comment>>,
    ) {
        for (line_idx, diff_line) in lines.iter().enumerate() {
            annotations.push(AnnotatedLine::DiffLine {
                file_idx,
                hunk_idx,
                line_idx,
                old_lineno: diff_line.old_lineno,
                new_lineno: diff_line.new_lineno,
            });

            // Line comments on old side (delete lines)
            if let Some(old_ln) = diff_line.old_lineno {
                Self::push_comments(
                    annotations,
                    file_idx,
                    Some(old_ln),
                    line_comments,
                    LineSide::Old,
                );
            }

            // Line comments on new side (added/context lines)
            if let Some(new_ln) = diff_line.new_lineno {
                Self::push_comments(
                    annotations,
                    file_idx,
                    Some(new_ln),
                    line_comments,
                    LineSide::New,
                );
            }
        }
    }

    /// Build annotations for side-by-side diff mode, pairing deletions and additions into aligned rows.
    fn build_side_by_side_annotations(
        annotations: &mut Vec<AnnotatedLine>,
        file_idx: usize,
        hunk_idx: usize,
        lines: &[travelagent_core::model::DiffLine],
        line_comments: &std::collections::HashMap<u32, Vec<travelagent_core::model::Comment>>,
    ) {
        let mut i = 0;
        while i < lines.len() {
            let diff_line = &lines[i];

            match diff_line.origin {
                LineOrigin::Context => {
                    annotations.push(AnnotatedLine::SideBySideLine {
                        file_idx,
                        hunk_idx,
                        del_line_idx: Some(i),
                        add_line_idx: Some(i),
                        old_lineno: diff_line.old_lineno,
                        new_lineno: diff_line.new_lineno,
                    });

                    // Push both Old- and New-side comments to match the SBS
                    // renderer's both-sides pass in
                    // `render_context_line_side_by_side` and the height
                    // estimator's post-PR-#113 both-sides count. Before this
                    // fix the annotation row count was shorter than the
                    // rendered row count whenever a context line had an
                    // Old-side comment, drifting cursor-to-row mapping and
                    // comment edit/delete by that offset.
                    Self::push_comments(
                        annotations,
                        file_idx,
                        diff_line.old_lineno,
                        line_comments,
                        LineSide::Old,
                    );
                    Self::push_comments(
                        annotations,
                        file_idx,
                        diff_line.new_lineno,
                        line_comments,
                        LineSide::New,
                    );

                    i += 1;
                }

                LineOrigin::Deletion => {
                    // Find consecutive deletions
                    let del_start = i;
                    let mut del_end = i + 1;
                    while del_end < lines.len() && lines[del_end].origin == LineOrigin::Deletion {
                        del_end += 1;
                    }

                    // Find consecutive additions following deletions
                    let add_start = del_end;
                    let mut add_end = add_start;
                    while add_end < lines.len() && lines[add_end].origin == LineOrigin::Addition {
                        add_end += 1;
                    }

                    let del_count = del_end - del_start;
                    let add_count = add_end - add_start;
                    let max_lines = del_count.max(add_count);

                    for offset in 0..max_lines {
                        let del_idx = if offset < del_count {
                            Some(del_start + offset)
                        } else {
                            None
                        };
                        let add_idx = if offset < add_count {
                            Some(add_start + offset)
                        } else {
                            None
                        };

                        let old_lineno = del_idx.and_then(|idx| lines[idx].old_lineno);
                        let new_lineno = add_idx.and_then(|idx| lines[idx].new_lineno);

                        annotations.push(AnnotatedLine::SideBySideLine {
                            file_idx,
                            hunk_idx,
                            del_line_idx: del_idx,
                            add_line_idx: add_idx,
                            old_lineno,
                            new_lineno,
                        });

                        Self::push_comments(
                            annotations,
                            file_idx,
                            old_lineno,
                            line_comments,
                            LineSide::Old,
                        );
                        Self::push_comments(
                            annotations,
                            file_idx,
                            new_lineno,
                            line_comments,
                            LineSide::New,
                        );
                    }

                    i = add_end;
                }
                LineOrigin::Addition => {
                    annotations.push(AnnotatedLine::SideBySideLine {
                        file_idx,
                        hunk_idx,
                        del_line_idx: None,
                        add_line_idx: Some(i),
                        old_lineno: None,
                        new_lineno: diff_line.new_lineno,
                    });

                    Self::push_comments(
                        annotations,
                        file_idx,
                        diff_line.new_lineno,
                        line_comments,
                        LineSide::New,
                    );

                    i += 1;
                }
            }
        }
    }

    /// What the cursor is on in a gap region
    pub fn get_gap_at_cursor(&self) -> Option<GapCursorHit> {
        let target = self.diff_state.cursor_line;
        match self.line_annotations.get(target) {
            Some(AnnotatedLine::Expander { gap_id, direction }) => {
                Some(GapCursorHit::Expander(gap_id.clone(), *direction))
            }
            Some(AnnotatedLine::HiddenLines { gap_id, .. }) => {
                Some(GapCursorHit::HiddenLines(gap_id.clone()))
            }
            Some(AnnotatedLine::ExpandedContext { gap_id, .. }) => {
                Some(GapCursorHit::ExpandedContent(gap_id.clone()))
            }
            _ => None,
        }
    }

    fn ensure_valid_tree_selection(&mut self) {
        let visible_items = self.build_visible_items();
        if visible_items.is_empty() {
            self.file_list_state.select(0);
            return;
        }

        let current_file_idx = self.diff_state.current_file_idx;
        let file_visible = visible_items.iter().any(|item| {
            matches!(item, FileTreeItem::File { file_idx, .. } if *file_idx == current_file_idx)
        });

        if file_visible {
            if let Some(tree_idx) = self.file_idx_to_tree_idx(current_file_idx) {
                self.file_list_state.select(tree_idx);
            }
        } else {
            if let Some(file) = self.diff_files.get(current_file_idx) {
                let file_path = file.display_path_lossy();
                let mut current = file_path.parent();
                while let Some(parent) = current {
                    if parent != Path::new("") {
                        let parent_str = parent.to_string_lossy().to_string();
                        for (tree_idx, item) in visible_items.iter().enumerate() {
                            if let FileTreeItem::Directory { path, .. } = item
                                && *path == parent_str
                            {
                                self.file_list_state.select(tree_idx);
                                return;
                            }
                        }
                    }
                    current = parent.parent();
                }
            }
            self.file_list_state.select(0);
        }
    }

    pub fn build_visible_items(&self) -> Vec<FileTreeItem> {
        let mut items = Vec::new();
        let mut seen_dirs: HashSet<String> = HashSet::new();

        for (file_idx, file) in self.diff_files.iter().enumerate() {
            let path = file.display_path_lossy();

            let mut ancestors: Vec<String> = Vec::new();
            let mut current = path.parent();
            while let Some(parent) = current {
                if parent != Path::new("") {
                    ancestors.push(parent.to_string_lossy().to_string());
                }
                current = parent.parent();
            }
            ancestors.reverse();

            let mut visible = true;
            for (depth, dir) in ancestors.iter().enumerate() {
                if !seen_dirs.contains(dir) && visible {
                    let expanded = self.ui_layout.expanded_dirs.contains(dir);
                    items.push(FileTreeItem::Directory {
                        path: dir.clone(),
                        depth,
                        expanded,
                    });
                    seen_dirs.insert(dir.clone());
                }

                if !self.ui_layout.expanded_dirs.contains(dir) {
                    visible = false;
                }
            }

            if visible {
                items.push(FileTreeItem::File {
                    file_idx,
                    depth: ancestors.len(),
                });
            }
        }

        items
    }

    pub fn get_selected_tree_item(&self) -> Option<FileTreeItem> {
        let visible_items = self.build_visible_items();
        let selected_idx = self.file_list_state.selected();
        visible_items.get(selected_idx).cloned()
    }
}