use std::ops;
use std::collections::HashMap;

use crate::parser::{self, GuardedBlock};
use parser::Atom;

pub struct Env<'a> {
    pub parent: Option<&'a Env<'a>>,
    pub values: HashMap<String, Atom>
}

impl<'a> Env<'a> {
    pub fn global() -> Self {
        Env {
            parent: None,
            values: HashMap::new(),
        }
    }
    pub fn child(parent: &'a Env<'a>) -> Self {
        Env {
            parent: Some(parent),
            values: HashMap::new(),
        }
    }
    pub fn lookup(&self, id: &str) -> Atom {
        if let Some(a) = self.values.get(id) {
            a.clone()
        } else if let Some(parent) = self.parent {
            parent.lookup(id)
        } else {
            panic!("Variable {} not in scope.", id)
        }
    }
    pub fn set(&mut self, id: String, atom: Atom) {
        self.values.insert(id, atom);
    }
}

pub fn eval(ast: &parser::Stmt, env: &mut Env) {
    match &ast {
        parser::Stmt::ReplPrint(expr) =>
            println!("{}", eval_expr(expr, env)),
        parser::Stmt::Assignment(id, expr) =>
            env.set(id.clone(), eval_expr(expr, env)),
        parser::Stmt::Conditional(guarded_blocks, default_block) => {
            let mut matched = false;
            for GuardedBlock { guard, block } in guarded_blocks {
                let res = eval_expr(&guard, env);
                match res {
                    Atom::Bool(true) => {
                        matched = true;
                        for stmt in block {
                            eval(stmt, env);
                        }
                        break;
                    }
                    Atom::Bool(false) => continue,
                    _ => panic!("Conditional expression does not evaluate to a bool"),
                }
            }
            if !matched {
                if let Some(block) = default_block {
                    for expr in default_block {
                        eval_expr(expr, env);
                    }
                }
            }
        }
        _ => todo!(),
    }
}

pub fn eval_expr(ast: &parser::Expr, env: &Env) -> Atom {
    match ast {
        parser::Expr::Id(a) => env.lookup(a),
        parser::Expr::Atom(a) => a.clone(),
        parser::Expr::UnaryMinus(a) => -eval_expr(a, env),
        parser::Expr::Plus(a, b) => eval_expr(a, env) + eval_expr(b, env),
        parser::Expr::Minus(a, b) => eval_expr(a, env) - eval_expr(b, env),
        parser::Expr::Mult(a, b) => eval_expr(a, env) * eval_expr(b, env),
        parser::Expr::Div(a, b) => eval_expr(a, env) / eval_expr(b, env),
        _ => panic!("Couldn't evalute expression {{ {:?} }}", ast),
    }
}

impl ops::Neg for Atom {
    type Output = Self;

    fn neg(self) -> Self {
        match self {
            Atom::Num(a) => Atom::Num(-a),
            _ => panic!("Can't negate non-numeral type!"),
        }
    }
}

impl ops::Add for Atom {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        match (self, other) {
            (Atom::Num(a), Atom::Num(b)) => Atom::Num(a + b),
            _ => panic!("Can't add non-numeral types!"),
        }
    }
}

impl ops::Sub for Atom {
    type Output = Self;

    fn sub(self, other: Self) -> Self {
        match (self, other) {
            (Atom::Num(a), Atom::Num(b)) => Atom::Num(a - b),
            _ => panic!("Can't subtract non-numeral types!"),
        }
    }
}

impl ops::Mul for Atom {
    type Output = Self;

    fn mul(self, other: Self) -> Self {
        match (self, other) {
            (Atom::Num(a), Atom::Num(b)) => Atom::Num(a * b),
            _ => panic!("Can't multiply non-numeral types!"),
        }
    }
}

impl ops::Div for Atom {
    type Output = Self;

    fn div(self, other: Self) -> Self {
        match (self, other) {
            (Atom::Num(a), Atom::Num(b)) => Atom::Num(a / b),
            _ => panic!("Can't divide non-numeral types!"),
        }
    }
}