//! Does whatever correctness and type checking necessary
//! to an AST.
//!
//! Okay, so here's what we need to do...
//! We can make this a one-pass compiler.  No recursion,
//! no forward decl.  Struct names are capital case,
//! variables are lower case -- does that matter?
//!
//! * Declare/define standard library types
//! * Enter a function:
//!    * Add variables to symbol tables -- can enumerate them in the process
//!    * Type check expressions
//!    * Make sure patterns/matches are valid
//!    * Whatever else the AST tells us to do
//! * Enter a struct definition and add it to the known types
//! * Handle entry point functions -- input and output types
//!   will need to go into SPIRV
//!
//! We can MOSTLY compile things one function at a time:
//! https://www.khronos.org/registry/spir-v/specs/1.0/SPIRV.html
//!
//! Each function can just be turned into a known-good AST
//! with whatever other info needs to be attached to it --
//! local symbol table for enumerated variables, etc.
//! Then we flatten it and break it into basic blocks
//! that do whatever SSA magic necessary.
//! Then we should just be able to emit code?
//!
//! Module also needs the entry point names and types,
//! the memory model (which can always be Logical addressing
//! and Simple memory model I think),
//! type decl's, constants, extended instruction sets
//! (which can include things like trig, etc),
//! debug info (OpName, OpLine, etc), capabilities,
//! whatever else.
//!
//! We're going to really want a good test program, and
//! a good set of SPIRV tools (disassembler, validator,
//! GLSL compiler to compare against).  And a debugger,
//! --renderdoc should be fine to start with.
//!
//! Bits that will be tricksy:
//!  * Structs, and associated pointers
//!  * Samplers and images

use std::cmp::Eq;
use std::collections::hash_map;
use std::fmt::Debug;
use std::hash::Hash;

// We use fnv so our hashmaps are always in deterministic
// order, so we have reproducable builds, so we can unit test
// things sanely.
use fnv::FnvHashMap;

use crate::ast::*;
use crate::Error;

/// A type definition, containing whatever we need
/// for that type.
///
/// Hmmm, how to represent vectors.  On the one hand,
/// they're just structs.  On the other hand, SPIR-V
/// has a particular type for vectors.  Hmm, for now
/// let's just treat them as structs, it might change
/// later.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum TypeDef {
    F32,
    Bool,
    Unit,
    Struct(Vec<(String, TypeDef)>),
    Function(Vec<TypeDef>, Box<TypeDef>),
}

/// A function definition; the AST and whatever
/// other information we care about that can get
/// fed to the actual compiler-y bits.
#[derive(Clone, Debug)]
pub struct FunctionDef {
    pub decl: FunctionDecl,
    pub functiontype: TypeDef,
}

/// Verification context.
/// Contains things like symbol tables, etc.
///
/// The `Default` impl will give you the default type
/// and function definitions, too.
#[derive(Clone, Debug)]
pub struct VContext {
    /// Function definitions.
    pub functions: FnvHashMap<String, FunctionDef>,
    /// Type definitions
    pub types: FnvHashMap<Type, TypeDef>,
}

impl Default for VContext {
    fn default() -> Self {
        let functions = FnvHashMap::default();
        let mut types = FnvHashMap::default();
        {
            // Default types.
            types.insert("F32".into(), TypeDef::F32);
            types.insert("Bool".into(), TypeDef::Bool);
            types.insert("()".into(), TypeDef::Unit);
            types.insert(
                "Vec4F".into(),
                TypeDef::Struct(vec![
                    ("x".into(), TypeDef::F32),
                    ("y".into(), TypeDef::F32),
                    ("z".into(), TypeDef::F32),
                    ("w".into(), TypeDef::F32),
                ]),
            );
        }
        Self { functions, types }
    }
}

fn hashtbl_insert_with_if_vacant<K, V>(
    tbl: &mut FnvHashMap<K, V>,
    name: K,
    vl: V,
) -> Result<(), Error>
where
    K: Debug + Hash + Eq,
{
    // TODO: Remove this redundant string alloc.
    let symbol_name = format!("{:?}", &name);
    match tbl.entry(name) {
        hash_map::Entry::Occupied(_) => Err(Error::SymbolExists(symbol_name.into())),
        hash_map::Entry::Vacant(e) => {
            e.insert(vl);
            Ok(())
        }
    }
}

impl VContext {
    fn function_exists(&self, name: &str) -> bool {
        self.functions.contains_key(name)
    }
    fn type_exists(&self, name: &str) -> bool {
        self.types.contains_key(&Type(name.to_string()))
    }

    /// Shortcut for testing a type mismatch.
    /// Basically a convenience for making a Result
    /// with the right error type.  Returns the found
    /// type if it matches the expected one.
    fn check_type_mismatch(&self, e: &Expr, expected: &Type) -> Result<Type, Error> {
        let tt = self.type_of_expr(e)?;
        if &tt != expected {
            Err(Error::TypeMismatch(expected.clone(), tt))
        } else {
            Ok(tt)
        }
    }

    /// Returns the type that the given expression must yield.
    /// Error if type is unknown, ie an undefined structure.
    pub fn type_of_expr(&self, e: &Expr) -> Result<Type, Error> {
        match e {
            Expr::Let(_, _, _) => Ok(Type("()".into())),
            Expr::Literal(l) => Ok(l.type_of()),
            Expr::If(test, ifpart, elsepart) => {
                let _ = self.check_type_mismatch(test, &Type("Bool".into()))?;
                let it = self.type_of_exprs(ifpart)?;
                let et = self.type_of_exprs(elsepart)?;
                if it != et {
                    Err(Error::TypeMismatch(it, et))
                } else {
                    Ok(it)
                }
            }
            Expr::Block(es) => self.type_of_exprs(es),
            _ => unimplemented!(),
        }
    }

    /// Returns the type that the given LIST of expressions must yield,
    /// which is the type of the last expression, or Unit if the
    /// list is empty.
    pub fn type_of_exprs(&self, e: &[Expr]) -> Result<Type, Error> {
        let e = e.last().unwrap_or(&Expr::Literal(Lit::Unit));
        self.type_of_expr(e)
    }

    /// Insert the given function into the function symbol table.
    /// Returns error if it already exists.
    fn define_function(&mut self, name: &str, def: FunctionDef) -> Result<(), Error> {
        hashtbl_insert_with_if_vacant(&mut self.functions, name.into(), def)
    }

    /// Insert the given type name into the type symbol table.
    /// Returns error if it already exists.
    fn define_type(&mut self, name: &str, def: TypeDef) -> Result<(), Error> {
        hashtbl_insert_with_if_vacant(&mut self.types, Type(name.into()), def)
    }

    /// Returns a TypeDef matching a given name.
    /// Assumes the typedef exists and has already been defined;
    /// may panic if it hasn't.
    ///
    /// It might be better to turn the AST into an IR with all types
    /// resolved, but, for now...
    pub fn get_defined_type(&self, name: &Type) -> &TypeDef {
        self.types.get(name).unwrap()
    }
}

// /// Takes a list of `Lit` and returns true if they're
// /// all the given type.
// pub fn lits_all_are_type(l: &[Lit], )

pub fn verify_expr(_ctx: &VContext, e: &Expr) -> Result<(), Error> {
    match e {
        Expr::BinOp(_op, _e1, _e2) => panic!("verify expr"),
        _ => unreachable!(),
    }
}

/// Validates whether a program is valid.
/// Returns the Context full of things ready to
/// compile.
pub fn verify(program: Vec<Decl>) -> Result<VContext, Error> {
    let mut ctx = VContext::default();
    for decl in program.iter() {
        match decl {
            Decl::Function(f) => {
                // define function type from args.
                let return_type = ctx.get_defined_type(&f.returns);
                let param_types: Vec<TypeDef> = f
                    .params
                    .iter()
                    .map(|p| ctx.get_defined_type(&p.typ).clone())
                    .collect();
                let functiontype = TypeDef::Function(param_types, Box::new(return_type.clone()));
                // TODO: This is stupid but easy, make it better sometime.
                let function_type_name = format!("{:?}", functiontype);
                // Functions may define the same type multiple times and it's not an error...
                if !ctx.type_exists(&function_type_name) {
                    ctx.define_type(&function_type_name, functiontype.clone())?;
                }
                // This clone may be expensive, but it only
                // happens once and if the definition fails
                // then it's a fatal error anyway, so.
                let def = FunctionDef {
                    decl: f.clone(),
                    functiontype,
                };
                ctx.define_function(f.name.as_str(), def)?;
            }
            Decl::Structure(_s) => {
                // ctx.define_type(name.as_ref(), StructureDecl)?;
            }
        }
    }

    verify_program(&mut ctx, &program)?;
    Ok(ctx)
}

/// Validate that whole-program properties are correct,
/// for instance that there are vertex and fragment functions
/// with reasonable types.
pub fn verify_program(ctx: &mut VContext, _program: &[Decl]) -> Result<(), Error> {
    if !ctx.function_exists("vertex") {
        Err(Error::Validation(
            "Required function `vertex` doesn't exist.".into(),
        ))
    } else if !ctx.function_exists("fragment") {
        Err(Error::Validation(
            "Required function `fragment` doesn't exist.".into(),
        ))
    } else {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::*;

    use crate::ast::*;
    use crate::verify;
    use crate::Error;
    #[test]
    fn required_functions_exist() {
        let vert = Decl::Function(FunctionDecl {
            name: String::from("vertex"),
            params: vec![],
            returns: Type("Vec4F".into()),
            body: vec![Expr::Structure(
                String::from("Vec4F"),
                vec![
                    (String::from("x"), Expr::Literal(Lit::F32(1.0))),
                    (String::from("y"), Expr::Literal(Lit::F32(1.0))),
                    (String::from("z"), Expr::Literal(Lit::F32(1.0))),
                    (String::from("w"), Expr::Literal(Lit::F32(1.0))),
                ],
            )],
        });
        let frag = Decl::Function(FunctionDecl {
            name: String::from("fragment"),
            params: vec![Param {
                name: "input".into(),
                typ: Type("Vec4F".into()),
            }],
            returns: Type("Vec4F".into()),
            body: vec![Expr::Var("input".into())],
        });

        // No functions
        assert_matches!(verify::verify(vec![]), Err(Error::Validation(_)));
        // only `fragment`
        assert_matches!(
            verify::verify(vec![frag.clone()]),
            Err(Error::Validation(_))
        );
        // only `vertex`
        assert_matches!(
            verify::verify(vec![vert.clone()]),
            Err(Error::Validation(_))
        );
        // both
        assert_matches!(verify::verify(vec![frag, vert]), Ok(_));
    }

    #[test]
    fn duplicate_functions_invalid() {
        let f = Decl::Function(FunctionDecl {
            name: String::from("foo"),
            params: vec![],
            returns: Type("Vec4F".into()),
            body: vec![Expr::Structure(
                String::from("Vec4F"),
                vec![
                    (String::from("x"), Expr::Literal(Lit::F32(1.0))),
                    (String::from("y"), Expr::Literal(Lit::F32(1.0))),
                    (String::from("z"), Expr::Literal(Lit::F32(1.0))),
                    (String::from("w"), Expr::Literal(Lit::F32(1.0))),
                ],
            )],
        });
        // Error here: vertex/fragment don't exist.  We need a test program that starts with
        // a minimal test program of some kind, maybe...
        //assert_matches!(verify::verify(vec![f.clone()]), Ok(_));
        assert_matches!(
            verify::verify(vec![f.clone(), f.clone()]),
            Err(Error::SymbolExists(_))
        );
    }

    fn test_type_of_exprs(es: &[(Expr, Result<Type, Error>)]) {
        let c = verify::VContext::default();
        for (e, t) in es {
            let computed_type = c.type_of_expr(&e);
            assert_eq!(t, &computed_type);
        }
    }

    #[test]
    fn test_type_of_exprs_lit() {
        let es = vec![
            (Expr::Literal(Lit::F32(1.0)), Ok(Type("F32".into()))),
            (Expr::Literal(Lit::Bool(false)), Ok(Type("Bool".into()))),
            (Expr::Literal(Lit::Unit), Ok(Type("()".into()))),
        ];
        test_type_of_exprs(&es);
    }

    #[test]
    fn test_type_of_exprs_if() {
        let es = vec![
            (
                // Test is wrong type
                Expr::If(
                    Box::new(Expr::Literal(Lit::F32(1.0))),
                    vec![Expr::Literal(Lit::F32(2.0))],
                    vec![Expr::Literal(Lit::F32(3.0))],
                ),
                Err(Error::TypeMismatch(Type("Bool".into()), Type("F32".into()))),
            ),
            (
                // Everything ok
                Expr::If(
                    Box::new(Expr::Literal(Lit::Bool(false))),
                    vec![Expr::Literal(Lit::F32(2.0))],
                    vec![Expr::Literal(Lit::F32(3.0))],
                ),
                Ok(Type("F32".into())),
            ),
            (
                // Test arms don't match type
                Expr::If(
                    Box::new(Expr::Literal(Lit::Bool(false))),
                    vec![Expr::Literal(Lit::F32(2.0))],
                    vec![Expr::Literal(Lit::Bool(true))],
                ),
                Err(Error::TypeMismatch(Type("F32".into()), Type("Bool".into()))),
            ),
        ];
        test_type_of_exprs(&es);
    }

    #[test]
    fn test_struct_structural_equality() {
        // So this is a little random.  But my plan was for two
        // struct's with identical fields to not have structural
        // equality in the LANGUAGE, but to indeed have structural
        // equality once actually compiled.  But they don't
        // seem to. So this is to test if they are actually Eq
        let def1 = verify::TypeDef::Struct(vec![
            ("thing1".into(), verify::TypeDef::F32),
            ("thing2".into(), verify::TypeDef::F32),
            ("thing3".into(), verify::TypeDef::Bool),
        ]);

        let def2 = verify::TypeDef::Struct(vec![
            ("thing1".into(), def1.clone()),
            ("thing2".into(), verify::TypeDef::F32),
            ("thing3".into(), def1.clone()),
        ]);
        let def3 = verify::TypeDef::Struct(vec![
            ("thing1".into(), def1.clone()),
            ("thing2".into(), verify::TypeDef::F32),
            ("thing3".into(), def1.clone()),
        ]);
        assert_eq!(&def2, &def3.clone());
        use std::hash::Hash;
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        assert_eq!(&def2.hash(&mut hasher), &def3.hash(&mut hasher));
        // ...hmmmm.  Everything seems fine.
    }
}