use std::rc::Rc;

use hashbrown::HashMap;

use nixjit_error::{Error, Result};
use nixjit_ir::{self as ir, ArgIdx, ExprId, PrimOpId};
use nixjit_lir as lir;
use nixjit_value::{Const, Symbol};

pub use crate::value::*;

mod value;

pub trait EvalContext: Sized {
    fn eval(&mut self, expr: &ExprId) -> Result<Value>;
    fn with_with_env<T>(
        &mut self,
        namespace: Rc<HashMap<String, Value>>,
        f: impl FnOnce(&mut Self) -> T,
    ) -> T;
    fn with_args_env<T>(&mut self, args: Vec<Value>, f: impl FnOnce(&mut Self) -> T) -> (Vec<Value>, T);
    fn lookup_with<'a>(&'a self, ident: &str) -> Option<&'a Value>;
    fn lookup_arg<'a>(&'a self, idx: ArgIdx) -> &'a Value;
    fn pop_frame(&mut self) -> Vec<Value>;
    fn call_primop(&mut self, id: PrimOpId, args: Vec<Value>) -> Result<Value>;
}

pub trait Evaluate<Ctx: EvalContext> {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value>;
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ExprId {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        ctx.eval(self)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for lir::Lir {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use lir::Lir::*;
        match self {
            AttrSet(x) => x.eval(ctx),
            List(x) => x.eval(ctx),
            HasAttr(x) => x.eval(ctx),
            BinOp(x) => x.eval(ctx),
            UnOp(x) => x.eval(ctx),
            Select(x) => x.eval(ctx),
            If(x) => x.eval(ctx),
            Call(x) => x.eval(ctx),
            With(x) => x.eval(ctx),
            Assert(x) => x.eval(ctx),
            ConcatStrings(x) => x.eval(ctx),
            Const(x) => x.eval(ctx),
            Str(x) => x.eval(ctx),
            Var(x) => x.eval(ctx),
            Path(x) => x.eval(ctx),
            ExprRef(expr) => ctx.eval(expr),
            FuncRef(func) => Ok(Value::Func(unsafe { func.clone() })),
            &ArgRef(idx) => Ok(ctx.lookup_arg(idx).clone()),
            &PrimOp(primop) => Ok(Value::PrimOp(primop)),
        }
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::AttrSet {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let mut attrs = AttrSet::new(
            self.stcs
                .iter()
                .map(|(k, v)| {
                    let eval_result = v.eval(ctx);
                    Ok((k.clone(), eval_result?))
                })
                .collect::<Result<_>>()?,
        );
        for (k, v) in self.dyns.iter() {
            let mut k = k.eval(ctx)?;
            k.coerce_to_string();
            let v_eval_result = v.eval(ctx)?;
            attrs.push_attr(k.unwrap_string(), v_eval_result);
        }
        let result = Value::AttrSet(attrs.into()).ok();
        Ok(result.unwrap())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::List {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let items = self
            .items
            .iter()
            .map(|val| val.eval(ctx))
            .collect::<Result<Vec<_>>>()?;
        let result = Value::List(List::from(items).into()).ok();
        Ok(result.unwrap())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::HasAttr {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use ir::Attr::*;
        let mut val = self.lhs.eval(ctx)?;
        val.has_attr(self.rhs.iter().map(|attr| {
            Ok(match attr {
                Str(ident) => ident.clone(),
                Dynamic(expr) => {
                    let mut val = expr.eval(ctx)?;
                    val.coerce_to_string();
                    val.unwrap_string()
                }
            })
        }))?;
        let result = val.ok();
        Ok(result.unwrap())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::BinOp {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use ir::BinOpKind::*;
        let mut lhs = self.lhs.eval(ctx)?;
        let mut rhs = self.rhs.eval(ctx)?;
        match self.kind {
            Add => lhs.add(rhs)?,
            Sub => {
                rhs.neg();
                lhs.add(rhs)?;
            }
            Mul => lhs.mul(rhs),
            Div => lhs.div(rhs)?,
            Eq => Value::eq(&mut lhs, &rhs),
            Neq => {
                Value::eq(&mut lhs, &rhs);
                lhs.not();
            }
            Lt => lhs.lt(rhs),
            Gt => {
                rhs.lt(lhs);
                lhs = rhs;
            }
            Leq => {
                rhs.lt(lhs);
                rhs.not();
                lhs = rhs;
            }
            Geq => {
                lhs.lt(rhs);
                lhs.not();
            }
            And => lhs.and(rhs),
            Or => lhs.or(rhs),
            Impl => {
                lhs.not();
                lhs.or(rhs);
            }
            Con => lhs.concat(rhs),
            Upd => lhs.update(rhs),
            PipeL => lhs.call(core::iter::once(Ok(rhs)), ctx)?,
            PipeR => {
                rhs.call(core::iter::once(Ok(lhs)), ctx)?;
                lhs = rhs;
            }
        }
        Ok(lhs)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::UnOp {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use ir::UnOpKind::*;
        let mut rhs = self.rhs.eval(ctx)?;
        match self.kind {
            Neg => {
                rhs.neg();
            }
            Not => {
                rhs.not();
            }
        };
        Ok(rhs)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Select {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use ir::Attr::*;
        let mut val = self.expr.eval(ctx)?;
        if let Some(default) = &self.default {
            let default = default.eval(ctx)?;
            val.select_with_default(
                self.attrpath.iter().map(|attr| {
                    Ok(match attr {
                        Str(ident) => ident.clone(),
                        Dynamic(expr) => {
                            let mut val = expr.eval(ctx)?;
                            val.coerce_to_string();
                            val.unwrap_string()
                        }
                    })
                }),
                default,
            )?;
        } else {
            val.select(self.attrpath.iter().map(|attr| {
                Ok(match attr {
                    Str(ident) => ident.clone(),
                    Dynamic(expr) => {
                        let mut val = expr.eval(ctx)?;
                        val.coerce_to_string();
                        val.unwrap_string()
                    }
                })
            }))?;
        }
        let result = val.ok();
        Ok(result.unwrap())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::If {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        // TODO: Error Handling
        let cond = self.cond.eval(ctx)?;
        let cond = cond
            .try_unwrap_bool()
            .map_err(|_| Error::EvalError(format!("expected a boolean but found ...")))?;

        if cond {
            self.consq.eval(ctx)
        } else {
            self.alter.eval(ctx)
        }
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Call {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let mut func = self.func.eval(ctx)?;
        // FIXME: ?
        let ctx_mut = unsafe { &mut *(ctx as *mut Ctx) };
        func.call(
            self.args
                .iter()
                .map(|arg| arg.eval(ctx)),
            ctx_mut,
        )?;
        Ok(func.ok().unwrap())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::With {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let namespace = self.namespace.eval(ctx)?;
        ctx.with_with_env(
            namespace
                .try_unwrap_attr_set()
                .map_err(|_| Error::EvalError(format!("expected a set but found ...")))?
                .into_inner(),
            |ctx| self.expr.eval(ctx),
        )
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Assert {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let cond = self.assertion.eval(ctx)?;
        let cond = cond
            .try_unwrap_bool()
            .map_err(|_| Error::EvalError(format!("expected a boolean but found ...")))?;
        if cond {
            self.expr.eval(ctx)
        } else {
            Err(Error::EvalError("assertion failed".to_string()))
        }
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::ConcatStrings {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let mut parts = self
            .parts
            .iter()
            .map(|part| {
                let mut part = part.eval(ctx)?;
                part.coerce_to_string();
                part.ok()
            })
            .collect::<Result<Vec<_>>>()?
            .into_iter();
        let init = parts.next().unwrap();
        let result = parts.fold(init, |mut a, b| {
            a.concat_string(b);
            a
        });
        Ok(result.ok().unwrap())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Str {
    fn eval(&self, _: &mut Ctx) -> Result<Value> {
        let result = Value::String(self.val.clone()).ok();
        Ok(result.unwrap())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Const {
    fn eval(&self, _: &mut Ctx) -> Result<Value> {
        let result = match self.val {
            Const::Null => Value::Null,
            Const::Int(x) => Value::Int(x),
            Const::Float(x) => Value::Float(x),
            Const::Bool(x) => Value::Bool(x),
        }
        .ok();
        Ok(result.unwrap())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Var {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        ctx.lookup_with(&self.sym)
            .ok_or_else(|| {
                Error::EvalError(format!(
                    "variable {} not found",
                    Symbol::from(self.sym.clone())
                ))
            })
            .map(|val| val.clone())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Path {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        todo!()
    }
}