//! Evaluation of syntax trees.

#[macro_use]
mod value;
mod capture;
mod ops;
mod scope;

pub use capture::*;
pub use scope::*;
pub use value::*;

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::rc::Rc;

use crate::cache::Cache;
use crate::color::Color;
use crate::diag::{Diag, DiagSet, Pass};
use crate::geom::{Angle, Length, Relative};
use crate::loading::{FileHash, Loader};
use crate::parse::parse;
use crate::syntax::visit::Visit;
use crate::syntax::*;

/// Evaluated a parsed source file into a module.
///
/// The `path` should point to the source file for the `tree` and is used to
/// resolve relative path names.
///
/// The `scope` consists of the base definitions that are present from the
/// beginning (typically, the standard library).
pub fn eval(
    loader: &mut dyn Loader,
    cache: &mut Cache,
    path: &Path,
    tree: Rc<Tree>,
    base: &Scope,
) -> Pass<Module> {
    let mut ctx = EvalContext::new(loader, cache, path, base);
    let map = tree.eval(&mut ctx);
    let module = Module {
        scope: ctx.scopes.top,
        template: vec![TemplateNode::Tree { tree, map }],
    };
    Pass::new(module, ctx.diags)
}

/// An evaluated module, ready for importing or execution.
#[derive(Debug, Clone, PartialEq)]
pub struct Module {
    /// The top-level definitions that were bound in this module.
    pub scope: Scope,
    /// The template defined by this module.
    pub template: TemplateValue,
}

/// The context for evaluation.
pub struct EvalContext<'a> {
    /// The loader from which resources (files and images) are loaded.
    pub loader: &'a mut dyn Loader,
    /// A cache for loaded resources.
    pub cache: &'a mut Cache,
    /// The active scopes.
    pub scopes: Scopes<'a>,
    /// Evaluation diagnostics.
    pub diags: DiagSet,
    /// The stack of imported files that led to evaluation of the current file.
    pub route: Vec<FileHash>,
    /// The location of the currently evaluated file.
    pub path: PathBuf,
    /// A map of loaded module.
    pub modules: HashMap<FileHash, Module>,
}

impl<'a> EvalContext<'a> {
    /// Create a new evaluation context with a base scope.
    pub fn new(
        loader: &'a mut dyn Loader,
        cache: &'a mut Cache,
        path: &Path,
        base: &'a Scope,
    ) -> Self {
        let mut route = vec![];
        if let Some(hash) = loader.resolve(path) {
            route.push(hash);
        }

        Self {
            loader,
            cache,
            scopes: Scopes::with_base(Some(base)),
            diags: DiagSet::new(),
            route,
            path: path.to_owned(),
            modules: HashMap::new(),
        }
    }

    /// Resolve a path relative to the current file.
    ///
    /// Generates an error if the file is not found.
    pub fn resolve(&mut self, path: &str, span: Span) -> Option<(PathBuf, FileHash)> {
        let dir = self.path.parent().expect("location is a file");
        let path = dir.join(path);
        match self.loader.resolve(&path) {
            Some(hash) => Some((path, hash)),
            None => {
                self.diag(error!(span, "file not found"));
                None
            }
        }
    }

    /// Process an import of a module relative to the current location.
    pub fn import(&mut self, path: &str, span: Span) -> Option<FileHash> {
        let (resolved, hash) = self.resolve(path, span)?;

        // Prevent cycling importing.
        if self.route.contains(&hash) {
            self.diag(error!(span, "cyclic import"));
            return None;
        }

        if self.modules.get(&hash).is_some() {
            return Some(hash);
        }

        let buffer = self.loader.load_file(&resolved).or_else(|| {
            self.diag(error!(span, "failed to load file"));
            None
        })?;

        let string = std::str::from_utf8(&buffer).ok().or_else(|| {
            self.diag(error!(span, "file is not valid utf-8"));
            None
        })?;

        // Prepare the new context.
        self.route.push(hash);
        let new_scopes = Scopes::with_base(self.scopes.base);
        let old_scopes = std::mem::replace(&mut self.scopes, new_scopes);

        // Evaluate the module.
        let tree = Rc::new(parse(string).output);
        let map = tree.eval(self);

        // Restore the old context.
        let new_scopes = std::mem::replace(&mut self.scopes, old_scopes);
        self.route.pop();

        self.modules.insert(hash, Module {
            scope: new_scopes.top,
            template: vec![TemplateNode::Tree { tree, map }],
        });

        Some(hash)
    }

    /// Add a diagnostic.
    pub fn diag(&mut self, diag: Diag) {
        self.diags.insert(diag);
    }

    /// Cast a value to a type and diagnose a possible error / warning.
    pub fn cast<T>(&mut self, value: Value, span: Span) -> Option<T>
    where
        T: Cast<Value>,
    {
        if value == Value::Error {
            return None;
        }

        match T::cast(value) {
            CastResult::Ok(t) => Some(t),
            CastResult::Warn(t, m) => {
                self.diag(warning!(span, "{}", m));
                Some(t)
            }
            CastResult::Err(value) => {
                self.diag(error!(
                    span,
                    "expected {}, found {}",
                    T::TYPE_NAME,
                    value.type_name(),
                ));
                None
            }
        }
    }
}

/// Evaluate an expression.
pub trait Eval {
    /// The output of evaluating the expression.
    type Output;

    /// Evaluate the expression to the output value.
    fn eval(&self, ctx: &mut EvalContext) -> Self::Output;
}

impl Eval for Tree {
    type Output = NodeMap;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let mut map = NodeMap::new();

        for node in self {
            let value = if let Some(call) = node.desugar() {
                call.eval(ctx)
            } else if let Node::Expr(expr) = node {
                expr.eval(ctx)
            } else {
                continue;
            };

            map.insert(node as *const _, value);
        }

        map
    }
}

impl Eval for Expr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        match *self {
            Self::None(_) => Value::None,
            Self::Bool(_, v) => Value::Bool(v),
            Self::Int(_, v) => Value::Int(v),
            Self::Float(_, v) => Value::Float(v),
            Self::Length(_, v, unit) => Value::Length(Length::with_unit(v, unit)),
            Self::Angle(_, v, unit) => Value::Angle(Angle::with_unit(v, unit)),
            Self::Percent(_, v) => Value::Relative(Relative::new(v / 100.0)),
            Self::Color(_, v) => Value::Color(Color::Rgba(v)),
            Self::Str(_, ref v) => Value::Str(v.clone()),
            Self::Ident(ref v) => match ctx.scopes.get(&v) {
                Some(slot) => slot.borrow().clone(),
                None => {
                    ctx.diag(error!(v.span, "unknown variable"));
                    Value::Error
                }
            },
            Self::Array(ref v) => Value::Array(v.eval(ctx)),
            Self::Dict(ref v) => Value::Dict(v.eval(ctx)),
            Self::Template(ref v) => Value::Template(vec![v.eval(ctx)]),
            Self::Group(ref v) => v.eval(ctx),
            Self::Block(ref v) => v.eval(ctx),
            Self::Call(ref v) => v.eval(ctx),
            Self::Closure(ref v) => v.eval(ctx),
            Self::Unary(ref v) => v.eval(ctx),
            Self::Binary(ref v) => v.eval(ctx),
            Self::Let(ref v) => v.eval(ctx),
            Self::If(ref v) => v.eval(ctx),
            Self::While(ref v) => v.eval(ctx),
            Self::For(ref v) => v.eval(ctx),
            Self::Import(ref v) => v.eval(ctx),
            Self::Include(ref v) => v.eval(ctx),
        }
    }
}

impl Eval for ArrayExpr {
    type Output = ArrayValue;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        self.items.iter().map(|expr| expr.eval(ctx)).collect()
    }
}

impl Eval for DictExpr {
    type Output = DictValue;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        self.items
            .iter()
            .map(|Named { name, expr }| (name.string.clone(), expr.eval(ctx)))
            .collect()
    }
}

impl Eval for TemplateExpr {
    type Output = TemplateNode;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let tree = Rc::clone(&self.tree);
        let map = self.tree.eval(ctx);
        TemplateNode::Tree { tree, map }
    }
}

impl Eval for GroupExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        self.expr.eval(ctx)
    }
}

impl Eval for BlockExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        if self.scoping {
            ctx.scopes.enter();
        }

        let mut output = Value::None;
        for expr in &self.exprs {
            output = expr.eval(ctx);
        }

        if self.scoping {
            ctx.scopes.exit();
        }

        output
    }
}

impl Eval for UnaryExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let value = self.expr.eval(ctx);
        if value == Value::Error {
            return Value::Error;
        }

        let ty = value.type_name();
        let out = match self.op {
            UnOp::Pos => ops::pos(value),
            UnOp::Neg => ops::neg(value),
            UnOp::Not => ops::not(value),
        };

        if out == Value::Error {
            ctx.diag(error!(
                self.span,
                "cannot apply '{}' to {}",
                self.op.as_str(),
                ty,
            ));
        }

        out
    }
}

impl Eval for BinaryExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        match self.op {
            BinOp::Add => self.apply(ctx, ops::add),
            BinOp::Sub => self.apply(ctx, ops::sub),
            BinOp::Mul => self.apply(ctx, ops::mul),
            BinOp::Div => self.apply(ctx, ops::div),
            BinOp::And => self.apply(ctx, ops::and),
            BinOp::Or => self.apply(ctx, ops::or),
            BinOp::Eq => self.apply(ctx, ops::eq),
            BinOp::Neq => self.apply(ctx, ops::neq),
            BinOp::Lt => self.apply(ctx, ops::lt),
            BinOp::Leq => self.apply(ctx, ops::leq),
            BinOp::Gt => self.apply(ctx, ops::gt),
            BinOp::Geq => self.apply(ctx, ops::geq),
            BinOp::Assign => self.assign(ctx, |_, b| b),
            BinOp::AddAssign => self.assign(ctx, ops::add),
            BinOp::SubAssign => self.assign(ctx, ops::sub),
            BinOp::MulAssign => self.assign(ctx, ops::mul),
            BinOp::DivAssign => self.assign(ctx, ops::div),
        }
    }
}

impl BinaryExpr {
    /// Apply a basic binary operation.
    fn apply<F>(&self, ctx: &mut EvalContext, op: F) -> Value
    where
        F: FnOnce(Value, Value) -> Value,
    {
        // Short-circuit boolean operations.
        let lhs = self.lhs.eval(ctx);
        match (self.op, &lhs) {
            (BinOp::And, Value::Bool(false)) => return lhs,
            (BinOp::Or, Value::Bool(true)) => return lhs,
            _ => {}
        }

        let rhs = self.rhs.eval(ctx);
        if lhs == Value::Error || rhs == Value::Error {
            return Value::Error;
        }

        // Save type names before we consume the values in case of error.
        let types = (lhs.type_name(), rhs.type_name());
        let out = op(lhs, rhs);
        if out == Value::Error {
            self.error(ctx, types);
        }

        out
    }

    /// Apply an assignment operation.
    fn assign<F>(&self, ctx: &mut EvalContext, op: F) -> Value
    where
        F: FnOnce(Value, Value) -> Value,
    {
        let slot = if let Expr::Ident(id) = self.lhs.as_ref() {
            match ctx.scopes.get(id) {
                Some(slot) => Rc::clone(slot),
                None => {
                    ctx.diag(error!(self.lhs.span(), "unknown variable"));
                    return Value::Error;
                }
            }
        } else {
            ctx.diag(error!(self.lhs.span(), "cannot assign to this expression"));
            return Value::Error;
        };

        let rhs = self.rhs.eval(ctx);
        let mut mutable = match slot.try_borrow_mut() {
            Ok(mutable) => mutable,
            Err(_) => {
                ctx.diag(error!(self.lhs.span(), "cannot assign to a constant"));
                return Value::Error;
            }
        };

        let lhs = std::mem::take(&mut *mutable);
        let types = (lhs.type_name(), rhs.type_name());
        *mutable = op(lhs, rhs);

        if *mutable == Value::Error {
            self.error(ctx, types);
            return Value::Error;
        }

        Value::None
    }

    fn error(&self, ctx: &mut EvalContext, (a, b): (&str, &str)) {
        ctx.diag(error!(self.span, "{}", match self.op {
            BinOp::Add => format!("cannot add {} and {}", a, b),
            BinOp::Sub => format!("cannot subtract {1} from {0}", a, b),
            BinOp::Mul => format!("cannot multiply {} with {}", a, b),
            BinOp::Div => format!("cannot divide {} by {}", a, b),
            _ => format!("cannot apply '{}' to {} and {}", self.op.as_str(), a, b),
        }));
    }
}

impl Eval for CallExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let callee = self.callee.eval(ctx);
        if let Some(func) = ctx.cast::<FuncValue>(callee, self.callee.span()) {
            let mut args = self.args.eval(ctx);
            let returned = func(ctx, &mut args);
            args.finish(ctx);
            returned
        } else {
            Value::Error
        }
    }
}

impl Eval for CallArgs {
    type Output = FuncArgs;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let items = self.items.iter().map(|arg| arg.eval(ctx)).collect();
        FuncArgs { span: self.span, items }
    }
}

impl Eval for CallArg {
    type Output = FuncArg;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        match self {
            Self::Pos(expr) => FuncArg {
                name: None,
                value: Spanned::new(expr.eval(ctx), expr.span()),
            },
            Self::Named(Named { name, expr }) => FuncArg {
                name: Some(Spanned::new(name.string.clone(), name.span)),
                value: Spanned::new(expr.eval(ctx), expr.span()),
            },
        }
    }
}

impl Eval for ClosureExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let params = Rc::clone(&self.params);
        let body = Rc::clone(&self.body);

        // Collect the captured variables.
        let captured = {
            let mut visitor = CapturesVisitor::new(&ctx.scopes);
            visitor.visit_closure(self);
            visitor.finish()
        };

        let name = self.name.as_ref().map(|id| id.to_string());
        Value::Func(FuncValue::new(name, move |ctx, args| {
            // Don't leak the scopes from the call site. Instead, we use the
            // scope of captured variables we collected earlier.
            let prev = std::mem::take(&mut ctx.scopes);
            ctx.scopes.top = captured.clone();

            for param in params.iter() {
                // Set the parameter to `none` if the argument is missing.
                let value =
                    args.eat_expect::<Value>(ctx, param.as_str()).unwrap_or_default();
                ctx.scopes.def_mut(param.as_str(), value);
            }

            let value = body.eval(ctx);
            ctx.scopes = prev;
            value
        }))
    }
}

impl Eval for LetExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let value = match &self.init {
            Some(expr) => expr.eval(ctx),
            None => Value::None,
        };
        ctx.scopes.def_mut(self.binding.as_str(), value);
        Value::None
    }
}

impl Eval for IfExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let condition = self.condition.eval(ctx);
        if let Some(condition) = ctx.cast(condition, self.condition.span()) {
            if condition {
                self.if_body.eval(ctx)
            } else if let Some(else_body) = &self.else_body {
                else_body.eval(ctx)
            } else {
                Value::None
            }
        } else {
            Value::Error
        }
    }
}

impl Eval for WhileExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let mut output = vec![];
        loop {
            let condition = self.condition.eval(ctx);
            if let Some(condition) = ctx.cast(condition, self.condition.span()) {
                if condition {
                    match self.body.eval(ctx) {
                        Value::Template(v) => output.extend(v),
                        Value::Str(v) => output.push(TemplateNode::Str(v)),
                        Value::Error => return Value::Error,
                        _ => {}
                    }
                } else {
                    return Value::Template(output);
                }
            } else {
                return Value::Error;
            }
        }
    }
}

impl Eval for ForExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        macro_rules! iter {
            (for ($($binding:ident => $value:ident),*) in $iter:expr) => {{
                let mut output = vec![];
                ctx.scopes.enter();

                #[allow(unused_parens)]
                for ($($value),*) in $iter {
                    $(ctx.scopes.def_mut($binding.as_str(), $value);)*

                    match self.body.eval(ctx) {
                        Value::Template(v) => output.extend(v),
                        Value::Str(v) => output.push(TemplateNode::Str(v)),
                        Value::Error => {
                            ctx.scopes.exit();
                            return Value::Error;
                        }
                        _ => {}
                    }
                }

                ctx.scopes.exit();
                Value::Template(output)
            }};
        }

        let iter = self.iter.eval(ctx);
        match (self.pattern.clone(), iter) {
            (ForPattern::Value(v), Value::Str(string)) => {
                iter!(for (v => value) in string.chars().map(|c| Value::Str(c.into())))
            }
            (ForPattern::Value(v), Value::Array(array)) => {
                iter!(for (v => value) in array.into_iter())
            }
            (ForPattern::KeyValue(i, v), Value::Array(array)) => {
                iter!(for (i => idx, v => value) in array.into_iter().enumerate())
            }
            (ForPattern::Value(v), Value::Dict(dict)) => {
                iter!(for (v => value) in dict.into_iter().map(|p| p.1))
            }
            (ForPattern::KeyValue(k, v), Value::Dict(dict)) => {
                iter!(for (k => key, v => value) in dict.into_iter())
            }

            (ForPattern::KeyValue(_, _), Value::Str(_)) => {
                ctx.diag(error!(self.pattern.span(), "mismatched pattern"));
                Value::Error
            }

            (_, iter) => {
                if iter != Value::Error {
                    ctx.diag(error!(
                        self.iter.span(),
                        "cannot loop over {}",
                        iter.type_name(),
                    ));
                }
                Value::Error
            }
        }
    }
}

impl Eval for ImportExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let span = self.path.span();
        let path = self.path.eval(ctx);

        if let Some(path) = ctx.cast::<String>(path, span) {
            if let Some(hash) = ctx.import(&path, span) {
                let mut module = &ctx.modules[&hash];
                match &self.imports {
                    Imports::Wildcard => {
                        for (var, slot) in module.scope.iter() {
                            let value = slot.borrow().clone();
                            ctx.scopes.def_mut(var, value);
                        }
                    }
                    Imports::Idents(idents) => {
                        for ident in idents {
                            if let Some(slot) = module.scope.get(&ident) {
                                let value = slot.borrow().clone();
                                ctx.scopes.def_mut(ident.as_str(), value);
                            } else {
                                ctx.diag(error!(ident.span, "unresolved import"));
                                module = &ctx.modules[&hash];
                            }
                        }
                    }
                }

                return Value::None;
            }
        }

        Value::Error
    }
}

impl Eval for IncludeExpr {
    type Output = Value;

    fn eval(&self, ctx: &mut EvalContext) -> Self::Output {
        let span = self.path.span();
        let path = self.path.eval(ctx);

        if let Some(path) = ctx.cast::<String>(path, span) {
            if let Some(hash) = ctx.import(&path, span) {
                return Value::Template(ctx.modules[&hash].template.clone());
            }
        }

        Value::Error
    }
}