// Copyright ⓒ 2017 David Li.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use ::ast::{self, WithLocation};
use super::environment;
use super::frame;
use super::ir;
use super::translate::{self, Level, Translate};
use super::types::{self, Ty};

#[derive(Debug)]
pub enum TypeError {
    Unimplemented,
    LengthMismatch {
        expected: usize,
        actual: usize,
    },
    Mismatch {
        expected: Ty,
        actual: Ty,
    },
    UnboundName,
}

type TypeEnvironment<'a> = environment::Environment<'a, String, Ty>;
pub type Result<T> = ::std::result::Result<T, WithLocation<TypeError>>;

pub fn translate<'a, F: frame::Frame>(
    translate: &mut Translate<F>,
    level: &mut Level<'a, F>,
    program: &ast::Program)
    -> Result<(translate::Expression, Ty)> {
    let mut venv = TypeEnvironment::new(None);
    let mut tenv = TypeEnvironment::new(None);
    tenv.add_binding("int".into(), Ty::Int);
    tenv.add_binding("string".into(), Ty::String);
    trans_exp(translate, level, &mut venv, &mut tenv, &*program.0)
}

macro_rules! err {
    ($exp: expr, $err: expr) => {
        Err(WithLocation::new($err, $exp.start, $exp.end))
    }
}

fn trans_ty<'a>(
    venv: &mut TypeEnvironment<'a>,
    tenv: &mut TypeEnvironment<'a>,
    ty: &WithLocation<ast::Ty>) -> Result<Ty> {
    match ty.value {
        ast::Ty::Name(ref name) => {
            tenv.lookup(&name)
                .map(|v| v.clone())
                .ok_or(WithLocation::new(TypeError::UnboundName, ty.start, ty.end))
        }
        ast::Ty::Array(ref inner_ty) => {
            trans_ty(venv, tenv, &inner_ty)
        },
        ast::Ty::Record(ref fields) => {
            let mut result = vec![];
            for field in fields {
                result.push(types::RecordField::new(
                    field.name.clone(),
                    trans_ty(venv, tenv, &field.ty)?));
            }
            Ok(Ty::Record(result))
        },
    }
}

fn trans_decl<'a, 'b, F: frame::Frame>(
    tr: &mut Translate<F>,
    level: &mut Level<'b, F>,
    venv: &mut TypeEnvironment<'a>,
    tenv: &mut TypeEnvironment<'a>,
    decl: &WithLocation<ast::Declaration>) -> Result<Ty> {
    match decl.declaration {
        ast::DeclarationBody::Fun { ref ty, ref params, ref body } => {
            let declared_ty = if let &Some(ref ty) = ty {
                Some(ast::WithLocation::new(trans_ty(venv, tenv, &ty)?, ty.start, ty.end))
            } else { None };

            let mut new_venv = TypeEnvironment::new(Some(venv));
            let mut new_tenv = TypeEnvironment::new(Some(tenv));
            let mut arg_types = vec![];
            for param in params.iter() {
                let arg_ty = trans_ty(&mut new_venv, &mut new_tenv, &param.ty)?;
                arg_types.push(arg_ty.clone());
                new_venv.add_binding(param.name.value.clone(), arg_ty);
            }

            let (body_exp, body_ty) =
                trans_exp(tr, level, &mut new_venv, &mut new_tenv, &*body)?;

            if let Some(decl_ty) = declared_ty {
                if &decl_ty.value != &body_ty {
                    return err!(decl_ty, TypeError::Mismatch {
                        expected: decl_ty.value,
                        actual: body_ty,
                    });
                }
            }

            Ok(Ty::Function(arg_types, Box::new(body_ty)))
        }
        ast::DeclarationBody::Ty { ref ty } => {
            trans_ty(venv, tenv, ty)
        }
        ast::DeclarationBody::Var { ref ty, ref value } => {
            let (var_exp, actual_ty) =
                trans_exp(tr, level, venv, tenv, &value)?;
            if let &Some(ref ty) = ty {
                let decl_ty = trans_ty(venv, tenv, &ty)?;
                if decl_ty != actual_ty {
                    return err!(ty, TypeError::Mismatch {
                        expected: decl_ty,
                        actual: actual_ty,
                    });
                }
            }
            Ok(actual_ty)
        }
    }
}

fn trans_exp<'a, 'b, F: frame::Frame>(
    tr: &mut Translate<F>,
    level: &mut Level<'b, F>,
    venv: &mut TypeEnvironment<'a>,
    tenv: &mut TypeEnvironment<'a>,
    exp: &WithLocation<ast::Expression>)
    -> Result<(translate::Expression, Ty)> {
    use ast::Expression::*;

    match &exp.value {
        &Let(ref decls, ref body) => {
            let mut new_venv = TypeEnvironment::new(Some(venv));
            let mut new_tenv = TypeEnvironment::new(Some(tenv));
            for decl in decls.iter() {
                let decl_ty = trans_decl(tr, level, &mut new_venv, &mut new_tenv, decl)?;
                match decl.declaration {
                    ast::DeclarationBody::Fun { .. } | ast::DeclarationBody::Var { .. } => {
                        new_venv.add_binding(decl.name.clone(), decl_ty);
                    }
                    ast::DeclarationBody::Ty { .. } => {
                        new_tenv.add_binding(decl.name.clone(), decl_ty);
                    }
                }
            }
            trans_exp(tr, level, &mut new_venv, &mut new_tenv, &*body)
        },
        &Call(ref name, ref args) => {
            let mut arg_types = vec![];
            for arg in args {
                arg_types.push(trans_exp(tr, level, venv, tenv, arg)?);
            }

            let fun_ty = venv.lookup(name)
                .ok_or(WithLocation::new(TypeError::UnboundName,
                                         exp.start, exp.end))?;
            match fun_ty {
                &Ty::Function(ref expected_args, ref result) => {
                    if expected_args.len() != arg_types.len() {
                        return err!(exp, TypeError::LengthMismatch {
                            expected: expected_args.len(),
                            actual: arg_types.len(),
                        });
                    }
                    for (i, (&(_, ref provided_ty), expected_ty)) in
                        arg_types.iter().zip(expected_args).enumerate() {
                        if provided_ty != expected_ty {
                            return err!(args[i], TypeError::Mismatch {
                                expected: expected_ty.clone(),
                                actual: provided_ty.clone(),
                            })
                        }
                    }
                    err!(exp, TypeError::Unimplemented)
                    // Ok((**result).clone())
                },
                otherwise => {
                    err!(exp, TypeError::Mismatch {
                        // TODO: better way to handle this
                        expected: Ty::Function(
                            arg_types.into_iter().map(|v| v.1).collect(),
                            Box::new(Ty::Nil)),
                        actual: otherwise.clone(),
                    })
                }
            }
        },
        &UnaryOp(ref op, ref operand) => {
            use ast::UnaryOp::*;
            let (operand_exp, operand_ty) = trans_exp(tr, level, venv, tenv, operand)?;
            match op {
                &Neg | &Pos => {
                    match operand_ty {
                        Ty::Int => {
                            err!(exp, TypeError::Unimplemented)
                            // Ok(Ty::Int)
                        }
                        other => {
                            err!(operand, TypeError::Mismatch {
                                expected: Ty::Int,
                                actual: other,
                            })
                        }
                    }
                }
                &Not => {
                    err!(exp, TypeError::Unimplemented)
                }
            }
        },
        &BinOp(ref op, ref left, ref right) => {
            use ast::BinOp::*;
            let (left_exp, left_ty) = trans_exp(tr, level, venv, tenv, left)?;
            let (right_exp, right_ty) = trans_exp(tr, level, venv, tenv, right)?;
            match op {
                &Add => {
                    match (left_ty, right_ty) {
                        (Ty::Int, Ty::Int) => {
                            Ok((tr.make_binop(ir::Binop::Plus, left_exp, right_exp),
                                Ty::Int))
                        }
                        (Ty::String, Ty::String) => {
                            err!(exp, TypeError::Unimplemented)
                            // Ok(Ty::String)
                        }
                        (Ty::Int, other) | (other, Ty::Int) => {
                            err!(right, TypeError::Mismatch {
                                expected: Ty::Int,
                                actual: other,
                            })
                        }
                        (Ty::String, other) | (other, Ty::String) => {
                            err!(right, TypeError::Mismatch {
                                expected: Ty::String,
                                actual: other,
                            })
                        }
                        _ => {
                            err!(exp, TypeError::Unimplemented)
                        }
                    }
                }
                _ => {
                    err!(exp, TypeError::Unimplemented)
                }
            }
        },
        &Number(n) => Ok((tr.make_num(n), Ty::Int)),
        &String(_) => err!(exp, TypeError::Unimplemented),
        // &String(_) => Ok(Ty::String),
        &Name(ref name) => {
            if let Some(ty) = venv.lookup(name) {
                err!(exp, TypeError::Unimplemented)
                // Ok(ty.clone())
            }
            else {
                err!(exp, TypeError::UnboundName)
            }
        },
        &Nil => {
            Err(WithLocation::new(TypeError::Unimplemented, exp.start, exp.end))
        },
    }
}