gitstack 5.3.0

//! Graph color management
//!
//! Color palette and color assignment algorithm with emphasis on visibility

use std::collections::{HashMap, VecDeque};

use ratatui::style::Color;

use crate::tui::theme;

/// Color index reserved for the main branch
pub const MAIN_BRANCH_COLOR: usize = 0;

/// Size of the color palette
pub const PALETTE_SIZE: usize = 12;

/// Color palette used in the graph (12 colors)
///
/// Color scheme designed for visibility:
/// - Main branch: Sky (blends with background while remaining distinguishable)
/// - Other branches: Rotate through highly saturated colors
/// - Light/dark pairs to clearly distinguish adjacent branches
const GRAPH_PALETTE: [Color; 12] = theme::GRAPH_PALETTE;

/// Get a Color value from a color index
pub fn get_graph_color(idx: usize) -> Color {
    GRAPH_PALETTE[idx % GRAPH_PALETTE.len()]
}

/// Context information for color assignment
#[derive(Debug, Clone)]
pub struct ColorContext {
    /// Lane (column) number
    pub lane: usize,
    /// Whether this is the main branch
    pub is_main_branch: bool,
    /// Hash of the fork point commit (branches sharing the same fork point get different colors)
    pub fork_point: Option<String>,
    /// Color of the parent commit (for continuity)
    pub parent_color: Option<usize>,
}

impl ColorContext {
    /// Create a new ColorContext
    pub fn new(lane: usize) -> Self {
        Self {
            lane,
            is_main_branch: false,
            fork_point: None,
            parent_color: None,
        }
    }

    /// Set the main branch flag
    pub fn with_main_branch(mut self, is_main: bool) -> Self {
        self.is_main_branch = is_main;
        self
    }

    /// Set the fork point
    pub fn with_fork_point(mut self, fork_point: Option<String>) -> Self {
        self.fork_point = fork_point;
        self
    }

    /// Set the parent color
    pub fn with_parent_color(mut self, parent_color: Option<usize>) -> Self {
        self.parent_color = parent_color;
        self
    }
}

/// Penalty-based color assignment manager
///
/// Inspired by keifu's algorithm, assigns colors considering the following penalties:
/// 1. Adjacent lane penalty: Avoid the same color as adjacent lanes
/// 2. History penalty: Avoid recently used colors
/// 3. Fork siblings penalty: Avoid the same color as branches sharing the same fork point
pub struct PenaltyBasedColorAssigner {
    /// Color assigned to each lane
    lane_colors: HashMap<usize, usize>,
    /// History of recently used colors
    color_history: VecDeque<usize>,
    /// Fork point hash -> list of colors in use
    fork_colors: HashMap<String, Vec<usize>>,
    /// History size
    history_size: usize,
}

impl PenaltyBasedColorAssigner {
    /// Create a new PenaltyBasedColorAssigner
    pub fn new() -> Self {
        Self {
            lane_colors: HashMap::new(),
            color_history: VecDeque::new(),
            fork_colors: HashMap::new(),
            history_size: 8,
        }
    }

    /// Get the color for the main branch
    pub fn main_color(&self) -> usize {
        MAIN_BRANCH_COLOR
    }

    /// Assign a color considering context
    pub fn assign_with_context(&mut self, ctx: &ColorContext) -> usize {
        // Main branch always gets color 0
        if ctx.is_main_branch {
            self.update_state(ctx.lane, MAIN_BRANCH_COLOR, ctx.fork_point.as_ref());
            return MAIN_BRANCH_COLOR;
        }

        // Initialize penalty array
        let mut penalties = [0.0f32; PALETTE_SIZE];

        // 1. Adjacent lane penalty (weight: 10.0)
        //    High penalty because the same color as adjacent lanes is hard to distinguish
        for neighbor in [ctx.lane.saturating_sub(1), ctx.lane + 1] {
            if neighbor != ctx.lane {
                if let Some(&color) = self.lane_colors.get(&neighbor) {
                    penalties[color] += 10.0;
                    // Light penalty for similar colors too
                    if let Some(similar) = self.similar_color(color) {
                        penalties[similar] += 5.0;
                    }
                }
            }
        }

        // 2. History penalty (weight: 3.0 * decay)
        //    Avoid recently used colors (vertical duplication avoidance)
        for (age, &color) in self.color_history.iter().enumerate() {
            let decay = 1.0 - (age as f32 / self.history_size as f32);
            penalties[color] += 3.0 * decay;
        }

        // 3. Fork siblings penalty (weight: 8.0)
        //    Branches derived from the same fork point should have different colors
        if let Some(ref fork) = ctx.fork_point {
            if let Some(siblings) = self.fork_colors.get(fork) {
                for &color in siblings {
                    penalties[color] += 8.0;
                }
            }
        }

        // 4. High penalty for main branch color (0)
        penalties[MAIN_BRANCH_COLOR] += 100.0;

        // Select the color with the minimum penalty
        let best = (1..PALETTE_SIZE)
            .min_by(|&a, &b| {
                penalties[a]
                    .partial_cmp(&penalties[b])
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .unwrap_or(1);

        // Update state
        self.update_state(ctx.lane, best, ctx.fork_point.as_ref());

        best
    }

    /// Update state
    fn update_state(&mut self, lane: usize, color: usize, fork_point: Option<&String>) {
        self.lane_colors.insert(lane, color);

        self.color_history.push_front(color);
        if self.color_history.len() > self.history_size {
            self.color_history.pop_back();
        }

        if let Some(fork) = fork_point {
            self.fork_colors
                .entry(fork.clone())
                .or_default()
                .push(color);
        }
    }

    /// Get similar color (light/dark variation of the same color family)
    fn similar_color(&self, color: usize) -> Option<usize> {
        // Palette composition (12 colors):
        // 0: Sky, 1: Blue, 2: Teal, 3: Mauve, 4: Peach, 5: Pink
        // 6: Green, 7: Lavender, 8: Maroon, 9: Yellow, 10: Sapphire, 11: Flamingo
        match color {
            0 => Some(10), // Sky <-> Sapphire (both blue family)
            10 => Some(0),
            1 => Some(7), // Blue <-> Lavender (both purple-blue family)
            7 => Some(1),
            2 => Some(6), // Teal <-> Green (both green family)
            6 => Some(2),
            3 => Some(5), // Mauve <-> Pink (both pink-purple family)
            5 => Some(3),
            4 => Some(8), // Peach <-> Maroon (both warm family)
            8 => Some(4),
            9 => Some(11), // Yellow <-> Flamingo (both warm-light family)
            11 => Some(9),
            _ => None,
        }
    }

    /// Get the color of a lane
    pub fn get_lane_color(&self, lane: usize) -> Option<usize> {
        self.lane_colors.get(&lane).copied()
    }

    /// Release a lane
    pub fn release_lane(&mut self, lane: usize) {
        self.lane_colors.remove(&lane);
    }
}

impl Default for PenaltyBasedColorAssigner {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_get_graph_color() {
        assert_eq!(get_graph_color(0), theme::SKY);
        assert_eq!(get_graph_color(1), theme::BLUE);
        assert_eq!(get_graph_color(12), theme::SKY); // Wraparound
    }

    // PenaltyBasedColorAssigner tests

    #[test]
    fn test_penalty_assigner_main_color() {
        let assigner = PenaltyBasedColorAssigner::new();
        assert_eq!(assigner.main_color(), MAIN_BRANCH_COLOR);
    }

    #[test]
    fn test_penalty_assigner_main_branch_always_color_0() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let ctx = ColorContext::new(0).with_main_branch(true);
        let color = assigner.assign_with_context(&ctx);
        assert_eq!(color, MAIN_BRANCH_COLOR);
    }

    #[test]
    fn test_penalty_assigner_skips_main_color() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let ctx = ColorContext::new(1);
        let color = assigner.assign_with_context(&ctx);
        assert_ne!(color, MAIN_BRANCH_COLOR);
    }

    #[test]
    fn test_penalty_assigner_adjacent_lanes_different_colors() {
        let mut assigner = PenaltyBasedColorAssigner::new();

        // Assign color to lane 0
        let ctx0 = ColorContext::new(0);
        let color0 = assigner.assign_with_context(&ctx0);

        // Lane 1 is adjacent, so it should get a different color
        let ctx1 = ColorContext::new(1);
        let color1 = assigner.assign_with_context(&ctx1);

        assert_ne!(
            color0, color1,
            "Adjacent lanes should have different colors"
        );
    }

    #[test]
    fn test_penalty_assigner_fork_siblings_different_colors() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let fork_point = "base_commit".to_string();

        // Branches derived from the same fork point
        let ctx1 = ColorContext::new(1).with_fork_point(Some(fork_point.clone()));
        let color1 = assigner.assign_with_context(&ctx1);

        let ctx2 = ColorContext::new(2).with_fork_point(Some(fork_point.clone()));
        let color2 = assigner.assign_with_context(&ctx2);

        let ctx3 = ColorContext::new(3).with_fork_point(Some(fork_point));
        let color3 = assigner.assign_with_context(&ctx3);

        assert_ne!(color1, color2, "Fork siblings should have different colors");
        assert_ne!(color2, color3, "Fork siblings should have different colors");
        assert_ne!(color1, color3, "Fork siblings should have different colors");
    }

    #[test]
    fn test_penalty_assigner_history_avoidance() {
        let mut assigner = PenaltyBasedColorAssigner::new();

        // Assign multiple colors consecutively
        let mut colors = Vec::new();
        for i in 0..5 {
            let ctx = ColorContext::new(i * 3); // Use distant lanes
            let color = assigner.assign_with_context(&ctx);
            colors.push(color);
        }

        // The first few colors should not overlap (due to history penalty)
        for i in 0..4 {
            for j in (i + 1)..5 {
                if i + 1 < j {
                    // Non-consecutive entries are less likely to duplicate
                    // (Complete duplication avoidance is not guaranteed, but checked as a tendency)
                }
            }
        }
        // At least not all the same color
        let unique_count = colors
            .iter()
            .collect::<std::collections::HashSet<_>>()
            .len();
        assert!(
            unique_count >= 3,
            "Should have color variation due to history"
        );
    }

    #[test]
    fn test_penalty_assigner_get_lane_color() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let ctx = ColorContext::new(5);
        let color = assigner.assign_with_context(&ctx);

        assert_eq!(assigner.get_lane_color(5), Some(color));
        assert_eq!(assigner.get_lane_color(0), None);
    }

    #[test]
    fn test_penalty_assigner_release_lane() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let ctx = ColorContext::new(3);
        let _color = assigner.assign_with_context(&ctx);

        assert!(assigner.get_lane_color(3).is_some());
        assigner.release_lane(3);
        assert!(assigner.get_lane_color(3).is_none());
    }

    #[test]
    fn test_color_context_builder() {
        let ctx = ColorContext::new(5)
            .with_main_branch(true)
            .with_fork_point(Some("abc123".to_string()))
            .with_parent_color(Some(3));

        assert_eq!(ctx.lane, 5);
        assert!(ctx.is_main_branch);
        assert_eq!(ctx.fork_point, Some("abc123".to_string()));
        assert_eq!(ctx.parent_color, Some(3));
    }

    #[test]
    fn test_penalty_assigner_many_lanes() {
        let mut assigner = PenaltyBasedColorAssigner::new();

        // Assign 16 or more lanes
        for i in 0..20 {
            let ctx = ColorContext::new(i);
            let color = assigner.assign_with_context(&ctx);
            // Color should be within palette size
            assert!(
                color < PALETTE_SIZE,
                "Color index should be within palette size"
            );
        }
    }
}