use std::collections::HashMap;

use crate::ast::{hash_ast, ExprAst};
use crate::compiler::{compile_ast, compile_to_scalar, CompiledExpr};
use crate::support::Arc;

/// In-memory content-addressed store keyed by AST hash.
#[derive(Default)]
pub struct ExprStore {
    compiled: HashMap<u64, Arc<CompiledExpr>>,
    asts: HashMap<u64, ExprAst>,
}

impl ExprStore {
    pub fn new() -> Self {
        Self::default()
    }

    /// Insert an AST, compile it once, and return its hash.
    /// If the hash already exists, the cached entry is reused.
    pub fn insert(&mut self, ast: ExprAst) -> Result<u64, String> {
        let hash = hash_ast(&ast);
        if !self.compiled.contains_key(&hash) {
            let compiled = compile_ast(&ast)?;
            self.compiled.insert(hash, Arc::new(compiled));
            self.asts.insert(hash, ast);
        }
        Ok(hash)
    }

    /// Fetch a compiled expression by hash.
    pub fn get(&self, hash: u64) -> Option<Arc<CompiledExpr>> {
        self.compiled.get(&hash).cloned()
    }

    /// Fetch the original AST by hash.
    pub fn get_ast(&self, hash: u64) -> Option<&ExprAst> {
        self.asts.get(&hash)
    }

    /// Evaluate a stored scalar expression for the provided input.
    pub fn eval_scalar(&self, hash: u64, input: f64) -> Result<f64, String> {
        let compiled = self
            .get(hash)
            .ok_or_else(|| format!("hash {:016x} not found", hash))?;

        match compiled.as_ref() {
            CompiledExpr::F64(expr) => Ok(expr.eval(input)),
            CompiledExpr::Pair(_) => Err("expression returns a pair, not f64".into()),
            CompiledExpr::PairToF64(_) => Err("expression expects a pair input, not f64".into()),
        }
    }

    /// Insert an AST that must evaluate to f64 -> f64, returning its hash and compiled form.
    pub fn insert_scalar(
        &mut self,
        ast: ExprAst,
    ) -> Result<
        (
            u64,
            Arc<dyn crate::dsl::Expr<f64, Out = f64> + Send + Sync + 'static>,
        ),
        String,
    > {
        let hash = self.insert(ast.clone())?;
        let expr = compile_to_scalar(&ast)?;
        Ok((hash, expr))
    }
}