use std::collections::BTreeMap;
use cord_trig::{TrigGraph, TrigOp};
use serde::{Deserialize, Serialize};

/// Canonical name for each CORDIC-consuming operation.
/// Non-CORDIC ops (Add, Sub, etc.) don't appear here.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum CordicOp {
    Sin,
    Cos,
    Tan,
    Asin,
    Acos,
    Atan,
    Sinh,
    Cosh,
    Tanh,
    Asinh,
    Acosh,
    Atanh,
    Sqrt,
    Exp,
    Ln,
    Hypot,
    Atan2,
}

impl CordicOp {
    pub fn name(&self) -> &'static str {
        match self {
            Self::Sin => "sin",
            Self::Cos => "cos",
            Self::Tan => "tan",
            Self::Asin => "asin",
            Self::Acos => "acos",
            Self::Atan => "atan",
            Self::Sinh => "sinh",
            Self::Cosh => "cosh",
            Self::Tanh => "tanh",
            Self::Asinh => "asinh",
            Self::Acosh => "acosh",
            Self::Atanh => "atanh",
            Self::Sqrt => "sqrt",
            Self::Exp => "exp",
            Self::Ln => "ln",
            Self::Hypot => "hypot",
            Self::Atan2 => "atan2",
        }
    }

    pub fn all() -> &'static [CordicOp] {
        &[
            Self::Sin, Self::Cos, Self::Tan,
            Self::Asin, Self::Acos, Self::Atan,
            Self::Sinh, Self::Cosh, Self::Tanh,
            Self::Asinh, Self::Acosh, Self::Atanh,
            Self::Sqrt, Self::Exp, Self::Ln,
            Self::Hypot, Self::Atan2,
        ]
    }

    pub fn from_name(s: &str) -> Option<Self> {
        Self::all().iter().find(|op| op.name() == s).copied()
    }
}

/// On-disk TOML representation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CordicConfigFile {
    #[serde(default)]
    pub defaults: Defaults,
    #[serde(default, rename = "override")]
    pub overrides: Vec<Override>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Defaults {
    #[serde(default = "default_word_bits")]
    pub word_bits: u8,
    #[serde(default = "default_f")]
    pub f: u8,
}

impl Default for Defaults {
    fn default() -> Self {
        Self {
            word_bits: default_word_bits(),
            f: default_f(),
        }
    }
}

fn default_word_bits() -> u8 { 32 }
fn default_f() -> u8 { 32 }

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Override {
    /// Regex pattern matching operation names (e.g. "sin|cos", "sinh|cosh|tanh")
    pub op: String,
    /// Iteration count override
    pub f: u8,
}

/// Resolved configuration: each operation has a concrete iteration count.
#[derive(Debug, Clone)]
pub struct CordicConfig {
    pub word_bits: u8,
    pub default_f: u8,
    /// Per-operation iteration counts. Only contains ops actually used.
    pub ops: BTreeMap<CordicOp, u8>,
}

impl Default for CordicConfig {
    fn default() -> Self {
        Self {
            word_bits: 32,
            default_f: 32,
            ops: BTreeMap::new(),
        }
    }
}

impl CordicConfig {
    /// Scan a TrigGraph and build a config with default f for every CORDIC op found.
    pub fn scan(graph: &TrigGraph) -> Self {
        let mut config = Self::default();
        for op in &graph.nodes {
            if let Some(cop) = trig_op_to_cordic(op) {
                config.ops.entry(cop).or_insert(config.default_f);
            }
        }
        config
    }

    /// Scan a TrigGraph and build config from a file representation.
    /// Overrides are applied in order; last match wins.
    pub fn from_file(graph: &TrigGraph, file: &CordicConfigFile) -> Self {
        let mut config = Self {
            word_bits: file.defaults.word_bits,
            default_f: file.defaults.f,
            ops: BTreeMap::new(),
        };

        // Populate with defaults for all ops found in graph
        for op in &graph.nodes {
            if let Some(cop) = trig_op_to_cordic(op) {
                config.ops.entry(cop).or_insert(config.default_f);
            }
        }

        // Apply overrides in order
        for ov in &file.overrides {
            if let Ok(re) = regex::Regex::new(&ov.op) {
                for (&cop, f_val) in config.ops.iter_mut() {
                    if re.is_match(cop.name()) {
                        *f_val = ov.f;
                    }
                }
            }
        }

        config
    }

    /// Parse from TOML string, resolving against the given graph.
    pub fn from_toml(graph: &TrigGraph, toml_str: &str) -> Result<Self, toml::de::Error> {
        let file: CordicConfigFile = toml::from_str(toml_str)?;
        Ok(Self::from_file(graph, &file))
    }

    /// Serialize the current config to a TOML file representation.
    pub fn to_file(&self) -> CordicConfigFile {
        // Group ops by f value to produce compact overrides
        let mut by_f: BTreeMap<u8, Vec<CordicOp>> = BTreeMap::new();
        for (&op, &f) in &self.ops {
            if f != self.default_f {
                by_f.entry(f).or_default().push(op);
            }
        }

        let overrides = by_f.into_iter().map(|(f, ops)| {
            let pattern = ops.iter()
                .map(|op| op.name())
                .collect::<Vec<_>>()
                .join("|");
            Override { op: pattern, f }
        }).collect();

        CordicConfigFile {
            defaults: Defaults {
                word_bits: self.word_bits,
                f: self.default_f,
            },
            overrides,
        }
    }

    /// Serialize to TOML string.
    pub fn to_toml(&self) -> String {
        let file = self.to_file();
        toml::to_string_pretty(&file).unwrap_or_default()
    }

    /// Get iteration count for a specific operation.
    pub fn iterations_for(&self, op: CordicOp) -> u8 {
        self.ops.get(&op).copied().unwrap_or(self.default_f)
    }

    /// Get the set of unique iteration counts used, for precomputing tables.
    pub fn unique_iteration_counts(&self) -> Vec<u8> {
        let mut counts: Vec<u8> = self.ops.values().copied().collect();
        counts.sort();
        counts.dedup();
        if counts.is_empty() {
            counts.push(self.default_f);
        }
        counts
    }
}

fn trig_op_to_cordic(op: &TrigOp) -> Option<CordicOp> {
    match op {
        TrigOp::Sin(_) => Some(CordicOp::Sin),
        TrigOp::Cos(_) => Some(CordicOp::Cos),
        TrigOp::Tan(_) => Some(CordicOp::Tan),
        TrigOp::Asin(_) => Some(CordicOp::Asin),
        TrigOp::Acos(_) => Some(CordicOp::Acos),
        TrigOp::Atan(_) => Some(CordicOp::Atan),
        TrigOp::Sinh(_) => Some(CordicOp::Sinh),
        TrigOp::Cosh(_) => Some(CordicOp::Cosh),
        TrigOp::Tanh(_) => Some(CordicOp::Tanh),
        TrigOp::Asinh(_) => Some(CordicOp::Asinh),
        TrigOp::Acosh(_) => Some(CordicOp::Acosh),
        TrigOp::Atanh(_) => Some(CordicOp::Atanh),
        TrigOp::Sqrt(_) => Some(CordicOp::Sqrt),
        TrigOp::Exp(_) => Some(CordicOp::Exp),
        TrigOp::Ln(_) => Some(CordicOp::Ln),
        TrigOp::Hypot(_, _) => Some(CordicOp::Hypot),
        TrigOp::Atan2(_, _) => Some(CordicOp::Atan2),
        _ => None,
    }
}

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

    #[test]
    fn scan_discovers_ops() {
        let mut g = TrigGraph::new();
        let angle = g.push(TrigOp::Const(1.0));
        let s = g.push(TrigOp::Sin(angle));
        let c = g.push(TrigOp::Cos(angle));
        let _ = g.push(TrigOp::Add(s, c));
        g.set_output(3);

        let config = CordicConfig::scan(&g);
        assert!(config.ops.contains_key(&CordicOp::Sin));
        assert!(config.ops.contains_key(&CordicOp::Cos));
        assert!(!config.ops.contains_key(&CordicOp::Tan));
        assert_eq!(config.ops.len(), 2);
    }

    #[test]
    fn override_applies() {
        let mut g = TrigGraph::new();
        let a = g.push(TrigOp::Const(1.0));
        let _ = g.push(TrigOp::Sin(a));
        let _ = g.push(TrigOp::Cos(a));
        let _ = g.push(TrigOp::Sinh(a));
        g.set_output(1);

        let file = CordicConfigFile {
            defaults: Defaults { word_bits: 32, f: 32 },
            overrides: vec![
                Override { op: "sin|cos".into(), f: 24 },
                Override { op: "sinh".into(), f: 16 },
            ],
        };

        let config = CordicConfig::from_file(&g, &file);
        assert_eq!(config.iterations_for(CordicOp::Sin), 24);
        assert_eq!(config.iterations_for(CordicOp::Cos), 24);
        assert_eq!(config.iterations_for(CordicOp::Sinh), 16);
    }

    #[test]
    fn toml_roundtrip() {
        let mut g = TrigGraph::new();
        let a = g.push(TrigOp::Const(1.0));
        let _ = g.push(TrigOp::Sin(a));
        let _ = g.push(TrigOp::Cos(a));
        g.set_output(1);

        let mut config = CordicConfig::scan(&g);
        config.ops.insert(CordicOp::Sin, 24);

        let toml_str = config.to_toml();
        let config2 = CordicConfig::from_toml(&g, &toml_str).unwrap();
        assert_eq!(config2.iterations_for(CordicOp::Sin), 24);
        assert_eq!(config2.iterations_for(CordicOp::Cos), 32);
    }
}
