use crate::eval::{EvalError, PrimitiveType, Value};
use crate::syntax::{ConstantType, Expression, Name, Program, Statement, TopLevel};
use crate::util::scoped_map::ScopedMap;
use internment::ArcIntern;
use std::collections::HashMap;
use std::str::FromStr;

impl Program {
    /// Evaluate the program, returning either an error or a pair of the final value
    /// produced and the output printed to the console.
    ///
    /// Doing this evaluation is particularly useful for testing, to ensure that if we
    /// modify a program in some way it does the same thing on both sides of the
    /// transformation. It's also sometimes just nice to know what a program will be
    /// doing.
    ///
    /// Note that the errors here are slightly more strict that we enforce at runtime.
    /// For example, we check for overflow and underflow errors during evaluation, and
    /// we don't check for those in the compiled code.
    pub fn eval(&self) -> Result<(Value<Expression>, String), EvalError<Expression>> {
        let mut env = ScopedMap::new();
        let mut stdout = String::new();
        let mut last_result = Value::Void;

        for stmt in self.items.iter() {
            match stmt {
                TopLevel::Function(name, arg_names, _, body) => {
                    last_result = Value::Closure(
                        name.clone().map(Name::intern),
                        env.clone(),
                        arg_names
                            .iter()
                            .cloned()
                            .map(|(x, _)| Name::intern(x))
                            .collect(),
                        body.clone(),
                    );
                    if let Some(name) = name {
                        env.insert(name.clone().intern(), last_result.clone());
                    }
                }

                TopLevel::Statement(stmt) => last_result = stmt.eval(&mut stdout, &mut env)?,

                TopLevel::Structure(_, _, _) => {
                    last_result = Value::Void;
                }
            }
        }

        Ok((last_result, stdout))
    }
}

impl Statement {
    fn eval(
        &self,
        stdout: &mut String,
        env: &mut ScopedMap<ArcIntern<String>, Value<Expression>>,
    ) -> Result<Value<Expression>, EvalError<Expression>> {
        match self {
            Statement::Binding(_, name, value) => {
                let actual_value = value.eval(stdout, env)?;
                env.insert(name.clone().intern(), actual_value);
                Ok(Value::Void)
            }

            Statement::Print(loc, name) => {
                let value = env
                    .get(&name.clone().intern())
                    .ok_or_else(|| EvalError::LookupFailed(loc.clone(), name.name.clone()))?;
                let value = if let Value::Number(x) = value {
                    Value::U64(*x)
                } else {
                    value.clone()
                };
                let line = format!("{} = {}\n", name, value);
                stdout.push_str(&line);
                Ok(Value::Void)
            }

            Statement::Expression(e) => e.eval(stdout, env),
        }
    }
}

impl Expression {
    fn eval(
        &self,
        stdout: &mut String,
        env: &mut ScopedMap<ArcIntern<String>, Value<Expression>>,
    ) -> Result<Value<Expression>, EvalError<Expression>> {
        match self {
            Expression::Value(_, v) => match v {
                super::Value::Number(_, ty, v) => match ty {
                    None => Ok(Value::Number(*v)),
                    // FIXME: make these types validate their input size
                    Some(ConstantType::Void) => Ok(Value::Void),
                    Some(ConstantType::I8) => Ok(Value::I8(*v as i8)),
                    Some(ConstantType::I16) => Ok(Value::I16(*v as i16)),
                    Some(ConstantType::I32) => Ok(Value::I32(*v as i32)),
                    Some(ConstantType::I64) => Ok(Value::I64(*v as i64)),
                    Some(ConstantType::U8) => Ok(Value::U8(*v as u8)),
                    Some(ConstantType::U16) => Ok(Value::U16(*v as u16)),
                    Some(ConstantType::U32) => Ok(Value::U32(*v as u32)),
                    Some(ConstantType::U64) => Ok(Value::U64(*v)),
                },
            },

            Expression::Constructor(_, on, fields) => {
                let mut map = HashMap::with_capacity(fields.len());

                for (k, v) in fields.iter() {
                    map.insert(k.clone().intern(), v.eval(stdout, env)?);
                }

                Ok(Value::Structure(Some(on.clone().intern()), map))
            }

            Expression::Reference(loc, n) => env
                .get(&ArcIntern::new(n.clone()))
                .ok_or_else(|| EvalError::LookupFailed(loc.clone(), n.clone()))
                .cloned(),

            Expression::FieldRef(loc, expr, field) => {
                let struck = expr.eval(stdout, env)?;

                if let Value::Structure(on, mut fields) = struck {
                    if let Some(value) = fields.remove(&field.clone().intern()) {
                        Ok(value)
                    } else {
                        Err(EvalError::BadFieldForStructure(
                            loc.clone(),
                            on,
                            field.clone().intern(),
                        ))
                    }
                } else {
                    Err(EvalError::NoFieldForValue(
                        loc.clone(),
                        struck,
                        field.clone().intern(),
                    ))
                }
            }

            Expression::Cast(_, target, expr) => {
                let target_type = PrimitiveType::from_str(target)?;
                let value = expr.eval(stdout, env)?;
                Ok(target_type.safe_cast(&value)?)
            }

            Expression::Primitive(_, op, args) => {
                let mut arg_values = Vec::with_capacity(args.len());

                for arg in args.iter() {
                    // yay, recursion! makes this pretty straightforward
                    arg_values.push(arg.eval(stdout, env)?);
                }

                Ok(Value::calculate(op, arg_values)?)
            }

            Expression::Call(loc, fun, args) => {
                let function = fun.eval(stdout, env)?;

                match function {
                    Value::Closure(name, mut closure_env, arguments, body) => {
                        if args.len() != arguments.len() {
                            return Err(EvalError::WrongArgCount(
                                loc.clone(),
                                name,
                                arguments.len(),
                                args.len(),
                            ));
                        }

                        closure_env.new_scope();
                        for (name, value) in arguments.into_iter().zip(args.iter()) {
                            let value = value.eval(stdout, env)?;
                            closure_env.insert(name, value);
                        }
                        let result = body.eval(stdout, &mut closure_env)?;
                        closure_env.release_scope();
                        Ok(result)
                    }
                    _ => Err(EvalError::NotAFunction(loc.clone(), function)),
                }
            }

            Expression::Block(_, stmts) => {
                let mut result = Value::Void;

                for stmt in stmts.iter() {
                    result = stmt.eval(stdout, env)?;
                }

                Ok(result)
            }
        }
    }
}

#[test]
fn two_plus_three() {
    let input = Program::parse(0, "x = 2 + 3; print x;").expect("parse works");
    let (_, output) = input.eval().expect("runs successfully");
    assert_eq!("x = 5u64\n", &output);
}

#[test]
fn lotsa_math() {
    let input = Program::parse(0, "x = 2 + 3 * 10 / 5 - 1; print x;").expect("parse works");
    let (_, output) = input.eval().expect("runs successfully");
    assert_eq!("x = 7u64\n", &output);
}