use crate::models::{SankeyFlow, SankeyLayout, SankeyNode};

/// Compute a Sankey diagram layout.
///
/// `income_sources`: (label, amount, (r, g, b)) for each income category
/// `expense_categories`: (label, amount, (r, g, b)) for each expense category
/// `total_height`: pixel height of the diagram area
pub fn compute_sankey_layout(
    income_sources: &[(String, f64, (f64, f64, f64))],
    expense_categories: &[(String, f64, (f64, f64, f64))],
    total_height: f64,
) -> SankeyLayout {
    let total_income: f64 = income_sources.iter().map(|(_, v, _)| v).sum();
    let total_expense: f64 = expense_categories.iter().map(|(_, v, _)| v).sum();
    let net = total_income - total_expense;
    let max_side = total_income.max(total_expense).max(1.0);

    let padding = 4.0;

    // Layout left (income) nodes
    let mut left_nodes = Vec::new();
    let mut y = 0.0;
    let income_count = income_sources.len().max(1);
    let total_padding_left = padding * (income_count.saturating_sub(1)) as f64;
    let available_left = total_height - total_padding_left;
    for (label, value, color) in income_sources {
        let h = (value / max_side) * available_left;
        left_nodes.push(SankeyNode {
            label: label.clone(),
            value: *value,
            color: *color,
            y,
            height: h,
        });
        y += h + padding;
    }

    // Layout right (expense) nodes
    let mut right_nodes = Vec::new();
    y = 0.0;
    let expense_count = expense_categories.len().max(1);
    let total_padding_right = padding * (expense_count.saturating_sub(1)) as f64;
    let available_right = total_height - total_padding_right;
    for (label, value, color) in expense_categories {
        let h = (value / max_side) * available_right;
        right_nodes.push(SankeyNode {
            label: label.clone(),
            value: *value,
            color: *color,
            y,
            height: h,
        });
        y += h + padding;
    }

    // Center node (net/available)
    let center_height = (total_income / max_side) * available_left;
    let center_y = 0.0;

    // Flows from income -> center
    let mut flows_in = Vec::new();
    let mut from_y_cursor = 0.0;
    let mut to_y_cursor = 0.0;
    for (i, node) in left_nodes.iter().enumerate() {
        let w = (node.value / max_side) * available_left;
        flows_in.push(SankeyFlow {
            from_idx: i,
            to_idx: 0,
            value: node.value,
            from_y: from_y_cursor,
            to_y: to_y_cursor,
            width: w,
        });
        from_y_cursor += w + padding;
        to_y_cursor += w;
    }

    // Flows from center -> expenses
    let mut flows_out = Vec::new();
    let mut from_y_cursor = 0.0;
    let mut to_y_cursor = 0.0;
    for (i, node) in right_nodes.iter().enumerate() {
        let w = (node.value / max_side) * available_right;
        flows_out.push(SankeyFlow {
            from_idx: 0,
            to_idx: i,
            value: node.value,
            from_y: from_y_cursor,
            to_y: to_y_cursor,
            width: w,
        });
        from_y_cursor += w;
        to_y_cursor += w + padding;
    }

    SankeyLayout {
        left_nodes,
        right_nodes,
        center_y,
        center_height,
        flows_in,
        flows_out,
        net,
    }
}

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

    #[test]
    fn test_basic_layout() {
        let income = vec![("Salary".into(), 5000.0, (0.0, 0.8, 0.0))];
        let expenses = vec![
            ("Rent".into(), 2000.0, (0.8, 0.0, 0.0)),
            ("Food".into(), 1000.0, (0.8, 0.4, 0.0)),
        ];
        let layout = compute_sankey_layout(&income, &expenses, 400.0);
        assert_eq!(layout.left_nodes.len(), 1);
        assert_eq!(layout.right_nodes.len(), 2);
        assert!((layout.net - 2000.0).abs() < 0.01);
        assert_eq!(layout.flows_in.len(), 1);
        assert_eq!(layout.flows_out.len(), 2);
    }

    #[test]
    fn test_empty_inputs() {
        let layout = compute_sankey_layout(&[], &[], 400.0);
        assert!(layout.left_nodes.is_empty());
        assert!(layout.right_nodes.is_empty());
        assert!((layout.net - 0.0).abs() < 0.01);
    }

    #[test]
    fn test_proportional_heights() {
        let income = vec![
            ("Salary".into(), 3000.0, (0.0, 0.8, 0.0)),
            ("Freelance".into(), 1000.0, (0.0, 0.6, 0.0)),
        ];
        let expenses = vec![("Rent".into(), 2000.0, (0.8, 0.0, 0.0))];
        let layout = compute_sankey_layout(&income, &expenses, 400.0);
        // Salary should be 3x the height of Freelance
        let salary_h = layout.left_nodes[0].height;
        let freelance_h = layout.left_nodes[1].height;
        assert!((salary_h / freelance_h - 3.0).abs() < 0.1);
    }
}
