syntax-workout-core 0.2.0

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::measure::{Measure, WeightUnit};
use crate::node::{Node, NodeKind};
use crate::intensity::Intensity;
use crate::visit::Visit;

/// Epley formula: 1RM = weight * (1 + reps / 30)
pub fn estimate_1rm(weight: f64, reps: u32) -> f64 {
    if reps <= 1 {
        weight
    } else {
        weight * (1.0 + reps as f64 / 30.0)
    }
}

/// Counts all Set (Leaf) nodes in the tree.
pub struct SetCounter(pub u32);

impl Visit for SetCounter {
    fn visit_leaf(&mut self, _measures: &[Measure], _intensity: Option<&Intensity>, _ancestors: &[&Node]) {
        self.0 += 1;
    }
}

/// Sums weight * reps across all sets, normalised to kg.
pub struct VolumeCalculator(pub f64);

impl Visit for VolumeCalculator {
    fn visit_leaf(&mut self, measures: &[Measure], _intensity: Option<&Intensity>, _ancestors: &[&Node]) {
        if let (Some(weight_kg), Some(reps)) = (extract_weight_kg(measures), extract_reps(measures)) {
            self.0 += weight_kg * reps as f64;
        }
    }
}

/// Sums all reps across all sets.
pub struct RepCounter(pub u32);

impl Visit for RepCounter {
    fn visit_leaf(&mut self, measures: &[Measure], _intensity: Option<&Intensity>, _ancestors: &[&Node]) {
        if let Some(reps) = extract_reps(measures) {
            self.0 += reps;
        }
    }
}

/// Volume for a specific exercise name only.
pub struct ExerciseVolumeCalculator {
    pub exercise_name: String,
    pub volume: f64,
}

impl Visit for ExerciseVolumeCalculator {
    fn visit_leaf(&mut self, measures: &[Measure], _intensity: Option<&Intensity>, ancestors: &[&Node]) {
        let parent_exercise = find_parent_exercise_name(ancestors);
        if parent_exercise == Some(self.exercise_name.as_str()) {
            if let (Some(weight_kg), Some(reps)) = (extract_weight_kg(measures), extract_reps(measures)) {
                self.volume += weight_kg * reps as f64;
            }
        }
    }
}

/// Combined stats result -- computed in a single tree walk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TreeStats {
    pub total_sets: u32,
    pub total_reps: u32,
    pub total_volume_kg: f64,
    pub exercises: Vec<ExerciseStats>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExerciseStats {
    pub name: String,
    pub sets: u32,
    pub reps: u32,
    pub volume_kg: f64,
    pub estimated_1rm: Option<f64>,
}

/// Single-pass visitor that collects all stats.
pub struct AllStatsVisitor {
    pub total_sets: u32,
    pub total_reps: u32,
    pub total_volume_kg: f64,
    /// exercise name -> accumulated stats
    exercises: HashMap<String, ExerciseAccum>,
    /// preserve insertion order
    exercise_order: Vec<String>,
}

struct ExerciseAccum {
    sets: u32,
    reps: u32,
    volume_kg: f64,
    best_1rm: Option<f64>,
}

impl AllStatsVisitor {
    pub fn new() -> Self {
        Self {
            total_sets: 0,
            total_reps: 0,
            total_volume_kg: 0.0,
            exercises: HashMap::new(),
            exercise_order: Vec::new(),
        }
    }

    pub fn into_tree_stats(self) -> TreeStats {
        let exercises = self
            .exercise_order
            .iter()
            .filter_map(|name| {
                self.exercises.get(name).map(|acc| ExerciseStats {
                    name: name.clone(),
                    sets: acc.sets,
                    reps: acc.reps,
                    volume_kg: acc.volume_kg,
                    estimated_1rm: acc.best_1rm,
                })
            })
            .collect();

        TreeStats {
            total_sets: self.total_sets,
            total_reps: self.total_reps,
            total_volume_kg: self.total_volume_kg,
            exercises,
        }
    }
}

impl Visit for AllStatsVisitor {
    fn visit_leaf(&mut self, measures: &[Measure], _intensity: Option<&Intensity>, ancestors: &[&Node]) {
        self.total_sets += 1;

        let reps = extract_reps(measures);
        let weight_kg = extract_weight_kg(measures);

        if let Some(r) = reps {
            self.total_reps += r;
        }

        let set_volume = match (weight_kg, reps) {
            (Some(w), Some(r)) => w * r as f64,
            _ => 0.0,
        };
        self.total_volume_kg += set_volume;

        let set_1rm = match (weight_kg, reps) {
            (Some(w), Some(r)) if w > 0.0 && r > 0 => Some(estimate_1rm(w, r)),
            _ => None,
        };

        if let Some(exercise_name) = find_parent_exercise_name(ancestors) {
            let name = exercise_name.to_string();
            if !self.exercises.contains_key(&name) {
                self.exercise_order.push(name.clone());
                self.exercises.insert(
                    name.clone(),
                    ExerciseAccum {
                        sets: 0,
                        reps: 0,
                        volume_kg: 0.0,
                        best_1rm: None,
                    },
                );
            }
            let acc = self.exercises.get_mut(&name).unwrap();
            acc.sets += 1;
            if let Some(r) = reps {
                acc.reps += r;
            }
            acc.volume_kg += set_volume;
            if let Some(rm) = set_1rm {
                acc.best_1rm = Some(match acc.best_1rm {
                    Some(prev) if prev >= rm => prev,
                    _ => rm,
                });
            }
        }
    }
}

/// Compute all stats for a workout tree in a single pass.
pub fn compute_tree_stats(root: &Node) -> TreeStats {
    let mut visitor = AllStatsVisitor::new();
    visitor.visit_tree(root);
    visitor.into_tree_stats()
}

// ── helpers ──────────────────────────────────────────────────────────

fn extract_weight_kg(measures: &[Measure]) -> Option<f64> {
    measures.iter().find_map(|m| match m {
        Measure::Weight { amount, unit } => {
            let kg = match unit {
                WeightUnit::Kg => *amount,
                WeightUnit::Lbs => *amount * 0.453592,
            };
            Some(kg)
        }
        _ => None,
    })
}

fn extract_reps(measures: &[Measure]) -> Option<u32> {
    measures.iter().find_map(|m| match m {
        Measure::Reps(r) => Some(*r),
        _ => None,
    })
}

fn find_parent_exercise_name<'a>(ancestors: &[&'a Node]) -> Option<&'a str> {
    ancestors
        .iter()
        .rev()
        .find(|n| matches!(n.kind, NodeKind::Exercise))
        .and_then(|n| n.name.as_deref())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::execution_mode::ExecutionMode;
    use crate::node::*;
    use std::collections::BTreeMap;

    #[test]
    fn epley_formula_basic() {
        let result = estimate_1rm(80.0, 5);
        // 80 * (1 + 5/30) = 80 * 7/6 = 93.333...
        assert!((result - 93.333).abs() < 0.01);
    }

    #[test]
    fn epley_formula_single_rep() {
        assert_eq!(estimate_1rm(100.0, 1), 100.0);
    }

    fn make_straight_sets() -> Node {
        let make_set = |id: &str, weight: f64, reps: u32| Node {
            id: NodeId::from_string(id),
            kind: NodeKind::Set,
            name: None,
            children: vec![],
            payload: NodePayload::Leaf {
                measures: vec![
                    Measure::Weight { amount: weight, unit: WeightUnit::Kg },
                    Measure::Reps(reps),
                ],
                intensity: None,
            },
            metadata: BTreeMap::new(),
        };

        let exercise = Node {
            id: NodeId::from_string("ex-bench"),
            kind: NodeKind::Exercise,
            name: Some("Bench Press".into()),
            children: vec![
                make_set("s1", 80.0, 5),
                make_set("s2", 80.0, 5),
                make_set("s3", 80.0, 4),
            ],
            payload: NodePayload::Exercise {
                measures: vec![],
                intensity: None,
                rest_seconds: Some(180.0),
            },
            metadata: BTreeMap::new(),
        };

        let block = Node {
            id: NodeId::from_string("blk-1"),
            kind: NodeKind::Block,
            name: None,
            children: vec![exercise],
            payload: NodePayload::Block {
                execution_mode: ExecutionMode::Sequential,
                rest_seconds: None,
            },
            metadata: BTreeMap::new(),
        };

        Node {
            id: NodeId::from_string("sess-1"),
            kind: NodeKind::Session,
            name: Some("Push Day A".into()),
            children: vec![block],
            payload: NodePayload::Temporal { rest_seconds: None },
            metadata: BTreeMap::new(),
        }
    }

    fn make_superset() -> Node {
        let make_set = |id: &str, weight: f64, reps: u32| Node {
            id: NodeId::from_string(id),
            kind: NodeKind::Set,
            name: None,
            children: vec![],
            payload: NodePayload::Leaf {
                measures: vec![
                    Measure::Weight { amount: weight, unit: WeightUnit::Kg },
                    Measure::Reps(reps),
                ],
                intensity: None,
            },
            metadata: BTreeMap::new(),
        };

        let ex_row = Node {
            id: NodeId::from_string("ex-row"),
            kind: NodeKind::Exercise,
            name: Some("DB Row".into()),
            children: vec![make_set("s1", 30.0, 10), make_set("s2", 30.0, 10)],
            payload: NodePayload::Exercise {
                measures: vec![],
                intensity: None,
                rest_seconds: None,
            },
            metadata: BTreeMap::new(),
        };

        let ex_press = Node {
            id: NodeId::from_string("ex-press"),
            kind: NodeKind::Exercise,
            name: Some("Incline Press".into()),
            children: vec![make_set("s3", 25.0, 12), make_set("s4", 25.0, 12)],
            payload: NodePayload::Exercise {
                measures: vec![],
                intensity: None,
                rest_seconds: None,
            },
            metadata: BTreeMap::new(),
        };

        let block = Node {
            id: NodeId::from_string("blk-a"),
            kind: NodeKind::Block,
            name: Some("Superset A".into()),
            children: vec![ex_row, ex_press],
            payload: NodePayload::Block {
                execution_mode: ExecutionMode::Parallel,
                rest_seconds: Some(90.0),
            },
            metadata: BTreeMap::new(),
        };

        Node {
            id: NodeId::from_string("sess-1"),
            kind: NodeKind::Session,
            name: Some("Upper Hypertrophy".into()),
            children: vec![block],
            payload: NodePayload::Temporal { rest_seconds: None },
            metadata: BTreeMap::new(),
        }
    }

    fn make_easy_run() -> Node {
        let set = Node {
            id: NodeId::from_string("s1"),
            kind: NodeKind::Set,
            name: None,
            children: vec![],
            payload: NodePayload::Leaf {
                measures: vec![
                    Measure::Distance {
                        amount: 5.0,
                        unit: crate::measure::DistanceUnit::Kilometers,
                    },
                    Measure::Duration { seconds: 1650.0 },
                ],
                intensity: None,
            },
            metadata: BTreeMap::new(),
        };

        let exercise = Node {
            id: NodeId::from_string("ex-run"),
            kind: NodeKind::Exercise,
            name: Some("Easy Run".into()),
            children: vec![set],
            payload: NodePayload::Exercise {
                measures: vec![],
                intensity: None,
                rest_seconds: None,
            },
            metadata: BTreeMap::new(),
        };

        let block = Node {
            id: NodeId::from_string("blk-1"),
            kind: NodeKind::Block,
            name: None,
            children: vec![exercise],
            payload: NodePayload::Block {
                execution_mode: ExecutionMode::Sequential,
                rest_seconds: None,
            },
            metadata: BTreeMap::new(),
        };

        Node {
            id: NodeId::from_string("sess-1"),
            kind: NodeKind::Session,
            name: Some("Easy Run".into()),
            children: vec![block],
            payload: NodePayload::Temporal { rest_seconds: None },
            metadata: BTreeMap::new(),
        }
    }

    #[test]
    fn straight_sets_stats() {
        let root = make_straight_sets();
        let stats = compute_tree_stats(&root);

        assert_eq!(stats.total_sets, 3);
        assert_eq!(stats.total_reps, 14); // 5+5+4
        // volume: 80*5 + 80*5 + 80*4 = 400+400+320 = 1120
        assert!((stats.total_volume_kg - 1120.0).abs() < 0.01);

        assert_eq!(stats.exercises.len(), 1);
        let bench = &stats.exercises[0];
        assert_eq!(bench.name, "Bench Press");
        assert_eq!(bench.sets, 3);
        assert_eq!(bench.reps, 14);
        assert!((bench.volume_kg - 1120.0).abs() < 0.01);
        // best 1RM from 80kg x 5 reps = 93.33
        assert!(bench.estimated_1rm.is_some());
        assert!((bench.estimated_1rm.unwrap() - 93.333).abs() < 0.01);
    }

    #[test]
    fn superset_stats() {
        let root = make_superset();
        let stats = compute_tree_stats(&root);

        assert_eq!(stats.total_sets, 4);
        assert_eq!(stats.total_reps, 44); // 10+10+12+12
        // volume: 30*10 + 30*10 + 25*12 + 25*12 = 300+300+300+300 = 1200
        assert!((stats.total_volume_kg - 1200.0).abs() < 0.01);

        assert_eq!(stats.exercises.len(), 2);

        let row = &stats.exercises[0];
        assert_eq!(row.name, "DB Row");
        assert_eq!(row.sets, 2);
        assert_eq!(row.reps, 20);
        assert!((row.volume_kg - 600.0).abs() < 0.01);

        let press = &stats.exercises[1];
        assert_eq!(press.name, "Incline Press");
        assert_eq!(press.sets, 2);
        assert_eq!(press.reps, 24);
        assert!((press.volume_kg - 600.0).abs() < 0.01);
    }

    #[test]
    fn endurance_has_zero_volume() {
        let root = make_easy_run();
        let stats = compute_tree_stats(&root);

        assert_eq!(stats.total_sets, 1);
        assert_eq!(stats.total_reps, 0);
        assert!((stats.total_volume_kg - 0.0).abs() < 0.01);
    }

    #[test]
    fn all_stats_visitor_matches_individual_visitors() {
        let root = make_straight_sets();

        // Individual visitors
        let mut set_counter = SetCounter(0);
        set_counter.visit_tree(&root);

        let mut vol_calc = VolumeCalculator(0.0);
        vol_calc.visit_tree(&root);

        let mut rep_counter = RepCounter(0);
        rep_counter.visit_tree(&root);

        // Combined visitor
        let stats = compute_tree_stats(&root);

        assert_eq!(stats.total_sets, set_counter.0);
        assert_eq!(stats.total_reps, rep_counter.0);
        assert!((stats.total_volume_kg - vol_calc.0).abs() < 0.001);
    }

    #[test]
    fn exercise_volume_calculator() {
        let root = make_superset();
        let mut calc = ExerciseVolumeCalculator {
            exercise_name: "DB Row".to_string(),
            volume: 0.0,
        };
        calc.visit_tree(&root);
        assert!((calc.volume - 600.0).abs() < 0.01);
    }

    #[test]
    fn lbs_to_kg_conversion() {
        let set = Node {
            id: NodeId::from_string("s1"),
            kind: NodeKind::Set,
            name: None,
            children: vec![],
            payload: NodePayload::Leaf {
                measures: vec![
                    Measure::Weight { amount: 135.0, unit: WeightUnit::Lbs },
                    Measure::Reps(5),
                ],
                intensity: None,
            },
            metadata: BTreeMap::new(),
        };

        let exercise = Node {
            id: NodeId::from_string("ex"),
            kind: NodeKind::Exercise,
            name: Some("Squat".into()),
            children: vec![set],
            payload: NodePayload::Exercise {
                measures: vec![],
                intensity: None,
                rest_seconds: None,
            },
            metadata: BTreeMap::new(),
        };

        let root = Node {
            id: NodeId::from_string("sess"),
            kind: NodeKind::Session,
            name: None,
            children: vec![exercise],
            payload: NodePayload::Temporal { rest_seconds: None },
            metadata: BTreeMap::new(),
        };

        let stats = compute_tree_stats(&root);
        // 135 lbs = 61.23492 kg, * 5 reps = 306.1746
        assert!((stats.total_volume_kg - 306.1746).abs() < 0.01);
    }
}