pub mod const_fold;
pub mod cost;
pub mod error;
pub mod extract;
pub mod lang;
pub mod restore;
pub mod rules;
pub mod strip;

use egg::Runner;

use crate::diag::Diag;
use crate::ir::TypedDag;
use crate::opt::const_fold::const_fold;
use crate::opt::extract::extract_best;
use crate::opt::lang::MathLang;
use crate::opt::restore::restore_types_with_info;
use crate::opt::rules::all_rules;
use crate::opt::strip::strip_types;

/// Optimization mode: conservative (IEEE 754–preserving) vs aggressive (e.g. constant-chain folding).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum OptMode {
    /// No reassociation; constant folding only when both operands are already constants.
    #[default]
    Conservative,
    /// Enables rewrites that may change float rounding (e.g. reassociate add/mul so constants can be folded).
    Aggressive,
}

/// Configuration for the optimization pipeline (const_fold + e-graph + restore).
#[derive(Debug, Clone)]
pub struct OptConfig {
    pub mode: OptMode,
    /// Fold constant subexpressions in add/mul chains (e.g. `(x+1)+2` → `x+3`).
    /// When `None`, follows mode: `true` if Aggressive, `false` if Conservative.
    pub const_chain_fold: Option<bool>,
}

impl Default for OptConfig {
    fn default() -> Self {
        Self {
            mode: OptMode::Conservative,
            const_chain_fold: None,
        }
    }
}

impl OptConfig {
    /// Whether to run aggressive rewrites (reassociation) and a second const_fold after restore.
    pub fn const_chain_fold_enabled(&self) -> bool {
        self.const_chain_fold
            .unwrap_or(self.mode == OptMode::Aggressive)
    }
}

#[derive(Debug, thiserror::Error)]
pub enum OptError {
    #[error("only single-output DAG is currently supported by optimize")]
    MultiOutputUnsupported,
    #[error("unsupported call `{name}` in opt conversion")]
    UnsupportedCall { name: String },
    #[error("invalid expr during conversion: {msg}")]
    InvalidExpr { msg: String },
}

#[derive(Debug, Clone)]
pub struct OptResult {
    pub dag: TypedDag,
    /// All diagnostics collected during optimisation (promotions, const-fold warnings, etc.).
    pub diags: Vec<Diag>,
}

/// Run optimization with default config (conservative mode).
pub fn optimize_with_info(dag: &TypedDag) -> Result<OptResult, OptError> {
    optimize_with_info_and_config(dag, &OptConfig::default())
}

/// Run optimization with the given config (mode and optional const_chain_fold).
pub fn optimize_with_info_and_config(
    dag: &TypedDag,
    config: &OptConfig,
) -> Result<OptResult, OptError> {
    if dag.outputs.len() != 1 {
        return Err(OptError::MultiOutputUnsupported);
    }

    let folded = const_fold(dag);
    let (expr, type_map) = strip_types(&folded.dag)?;

    let rules = all_rules(config);
    let runner = Runner::<MathLang, (), ()>::default()
        .with_expr(&expr)
        .with_iter_limit(8)
        .with_node_limit(20_000)
        .run(&rules);

    let root = runner
        .roots
        .first()
        .copied()
        .ok_or_else(|| OptError::InvalidExpr {
            msg: "runner produced no roots".to_string(),
        })?;
    let (_, best_expr) = extract_best(&runner.egraph, root);
    let (mut restored, restore_diags) =
        restore_types_with_info(&best_expr, &type_map, &folded.dag.params)?;

    if config.const_chain_fold_enabled() {
        let refolded = const_fold(&restored);
        restored = refolded.dag;
    }

    let mut diags = folded.diags;
    diags.extend(restore_diags);

    Ok(OptResult {
        dag: restored,
        diags,
    })
}

pub fn optimize(dag: &TypedDag) -> Result<TypedDag, OptError> {
    optimize_with_info(dag).map(|r| r.dag)
}
