//! The actual compiler backend!  Should take in a VALIDATED AST and
//! turn it to SPIR-V.  Uses the `rspirv` crate for output.

use fnv::FnvHashMap;
use rspirv::mr::Builder;
use spirv_headers as spirv;

use crate::ast;
use crate::verify;

#[cfg(test)]
mod tests;

/// All the possible info attached to variable.
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct VarBinding {
    /// The variable's address word.
    pub address: spirv::Word,
    /// the variable's type
    pub typedef: spirv::Word,
}

/// Compilation context, containing the SPIRV builder and
/// such.
pub struct CContext {
    /// SPIR-V builder.  Sticking something in it involves no
    /// caching or anything, so to get rid of duplicate decl's
    /// we also have...
    pub b: Builder,
    /// A type table, that maps type definitions to their
    /// decl's in the builder.
    pub typetable: FnvHashMap<verify::TypeDef, spirv::Word>,
    /// Constants, maps literal values to declared words.
    pub consts: FnvHashMap<ast::Lit, spirv::Word>,
    /// Here's the actual scoping mechanism, a stack of
    /// name -> Word associations that go from variable
    /// or parameter names to the associated Word.
    pub symtable: Vec<FnvHashMap<String, VarBinding>>,
}

impl CContext {
    pub fn new() -> Self {
        let mut b = Builder::new();
        // Set up module stuff.
        b.set_version(1, 0);
        b.capability(spirv::Capability::Shader);
        b.memory_model(spirv::AddressingModel::Logical, spirv::MemoryModel::Simple);
        let typetable = FnvHashMap::default();
        let consts = FnvHashMap::default();
        // Global symbol table for functions, consts, etc...
        let symtable = vec![FnvHashMap::default()];
        Self {
            b,
            typetable,
            consts,
            symtable,
        }
    }

    /// Pushes a new scope to the symbol table.
    pub fn push_scope(&mut self) {
        self.symtable.push(Default::default());
    }

    /// Binds the variable to the current scope.
    /// Will overwrite any previous binding.  That's fine.
    ///
    /// Panics if no scope.
    pub fn bind_var(&mut self, name: &str, address: spirv::Word, typedef: spirv::Word) {
        self.symtable
            .last_mut()
            .expect("No scope in variable binding!")
            .insert(name.into(), VarBinding { address, typedef });
    }

    /// Returns the Word bound to the given var name, or panics if not found.
    /// All unknown vars should get caught by the validation step.
    pub fn lookup_var(&self, name: &str) -> &VarBinding {
        assert!(self.symtable.len() > 1, "No scope for variable lookup!");
        for scope in self.symtable.iter().rev() {
            if let Some(w) = scope.get(name) {
                return w;
            }
        }
        dbg!(&self.symtable);
        panic!("Variable not found!")
    }

    /// Pops the top scope from the symbol table.
    /// Panics on underflow.
    ///
    /// TODO: Just use Rust's scoping and destruction to do this?
    /// Meh.
    pub fn pop_scope(&mut self) {
        self.symtable
            .pop()
            .expect("Tried to pop empty scope stack!");
    }

    /// Defines a constant for a literal value
    pub fn define_const(&mut self, vl: ast::Lit) -> spirv::Word {
        // This is silly but doing the check inside the closure
        // causes borrowing issues.  :|
        let type_float = self.get_type(&verify::TypeDef::F32);
        let type_bool = self.get_type(&verify::TypeDef::Bool);
        //let type_unit = self.get_type(&verify::TypeDef::Unit);

        let consts = &mut self.consts;
        let b = &mut self.b;
        *consts.entry(vl.clone()).or_insert_with(|| match vl {
            ast::Lit::F32(f) => b.constant_f32(type_float, f),
            ast::Lit::Bool(bl) => {
                if bl {
                    b.constant_true(type_bool)
                } else {
                    b.constant_false(type_bool)
                }
            }
            // TODO: This is WRONG WRONG WRONG but
            // I'm not entirely sure how to make it right
            // yet.  Unit is not actually a value, so.
            ast::Lit::Unit =>
                0
                //b.constant_null(type_unit),
        })
    }

    /// Panics if the type does not exist.
    ///
    /// TODO: Should this take a name and a VContext instead of a TypeDef?
    pub fn get_type(&self, t: &verify::TypeDef) -> spirv::Word {
        *self.typetable.get(t).expect("Could not get type!")
    }

    /// Add a type to the type table, recursively
    /// if necessary.  Will also actually build the type
    /// into the SPIR-V builder.
    ///
    /// Returns the `Word` that is a handle to the defined type.
    /// If the type already exists, it will not duplicate it, just
    /// return the previous word.
    ///
    /// We do not track the types' names here, just the signatures,
    /// so two identical type defs will get coalesced to one.
    /// This may or may not be a good idea, debugging might want separate
    /// types referring to the same name.
    ///
    /// If we DO want to do that (and I'm leaning towards yes) we need to
    /// add the various OpName instructions here.
    pub fn add_type(&mut self, typedef: &verify::TypeDef) -> spirv::Word {
        use verify::TypeDef;
        match typedef {
            TypeDef::F32 => {
                // This helps rustc figure out a double-borrow on self
                let b = &mut self.b;
                *self
                    .typetable
                    .entry(typedef.clone())
                    .or_insert_with(|| b.type_float(32))
            }
            TypeDef::Bool => {
                let b = &mut self.b;
                *self
                    .typetable
                    .entry(typedef.clone())
                    .or_insert_with(|| b.type_bool())
            }
            TypeDef::Unit => {
                let b = &mut self.b;
                *self
                    .typetable
                    .entry(typedef.clone())
                    .or_insert_with(|| b.type_void())
            }
            TypeDef::Struct(fields) => {
                // Okay, since we recurse I guess we CAN'T use the Entry
                // API here.  Bah.  Bah, I say!
                // Okay we can but it's more of a pain.  Whatever.
                if let Some(val) = self.typetable.get(typedef) {
                    // Struct already exists
                    *val
                } else {
                    // println!("Adding new struct with fields: {:#?}", fields);
                    let mut fields_words = vec![];
                    for (_name, typedef) in fields.iter() {
                        // First we make sure all the types
                        // used by this struct
                        let type_word = self.add_type(typedef);
                        fields_words.push(type_word);
                    }
                    // Then make the struct itself
                    let n = self.b.type_struct(fields_words);
                    self.typetable.insert(typedef.clone(), n);
                    n
                }
            }
            TypeDef::Function(params, returns) => {
                if let Some(val) = self.typetable.get(typedef) {
                    // Function already exists
                    *val
                } else {
                    let rettype_word = self.add_type(returns);
                    let mut param_words = vec![];
                    for typedef in params.iter() {
                        let type_word = self.add_type(typedef);
                        param_words.push(type_word);
                    }
                    let f = self.b.type_function(rettype_word, param_words);
                    self.typetable.insert(typedef.clone(), f);
                    f
                }
            }
        }
    }

    /// Returns the Word holding the result fo the expression.
    pub fn compile_expr(
        &mut self,
        e: &ast::Expr,
        ctx: &verify::VContext,
    ) -> Result<spirv::Word, crate::Error> {
        let result_word = match e {
            ast::Expr::Var(name) => {
                let binding = self.lookup_var(name);
                let addr = binding.address;
                //let vartype = binding.typedef;
                addr
            }
            ast::Expr::Let(pattern, typ, vl) => {
                // Evaluate the expression
                let value_word = self.compile_expr(vl, ctx)?;
                // Find the type of the variable
                let var_typedef = ctx.get_defined_type(&typ);
                let var_typeword = self.get_type(var_typedef);
                // Allocate space for a variable
                //
                // TODO: I AM NOT CONVINCED that this is necessary,
                // since our values are all immutable!  We might be
                // able to just store the result value word of an
                // expression directly into our scope block and not
                // actually explicitly allocate local vars.
                // But glslang does it this way so for now we follow
                // its example.
                /*
                let var_addr =
                    self.b
                        .variable(var_typeword, None, spirv::StorageClass::Function, None);
                // Save the result of the expression to the var.
                self.b.store(var_addr, value_word, None, [])?;
                 */
                // TODO: For now, we don't actually implement any
                // pattern matching.
                match pattern {
                    ast::Pattern::FreeVar(varname) => {
                        self.bind_var(varname, value_word, var_typeword)
                    }
                }
                value_word
            }
            ast::Expr::Literal(lit) => {
                // TODO: DOUBLE CHECK: I think this returns the address
                // of the const, not the value
                self.define_const(lit.clone())
            }
            ast::Expr::FunCall(fname, params) => {
                // Compile arguments
                let param_words: Result<Vec<spirv::Word>, _> = params
                    .clone()
                    .iter()
                    .map(|param| self.compile_expr(param, ctx))
                    .collect();
                let param_words = param_words?;
                // Look up function word
                // ...which requires having compiled the function first.  :|
                // TODO: Fix!  Somehow...  If it isn't, this will panic
                // with variable not found.
                let binding = *self.lookup_var(&fname);
                // Get the function's return type.  We have to fetch this
                // from the VContext, since we can't really dig it back
                // out of the SPIR-V.
                let functiondef = ctx
                    .functions
                    .get(fname)
                    .expect("Function does not exist for funcall");
                if let verify::TypeDef::Function(ref _params, ref rettype) =
                    functiondef.functiontype
                {
                    // TODO: Tail call optimization would go here if it went anywhere.
                    let rettype_word = self.get_type(rettype);
                    self.b
                        .function_call(rettype_word, None, binding.address, param_words)?
                } else {
                    unreachable!("Function type is not TypeDef::Function")
                }
            }
            ast::Expr::Block(exprs) => {
                assert!(exprs.len() > 0, "Blocks with no expressions are verboten!");
                self.push_scope();
                let last_word: spirv::Word = exprs
                    .iter()
                    .map(|e| self.compile_expr(e, ctx))
                    .last()
                    .expect("Can't happen")?;
                self.pop_scope();
                last_word
            }
            ast::Expr::BinOp(op, e1, e2) => {
                let t1 = ctx.get_defined_type(&ctx.type_of_expr(e1)?);
                let t2 = ctx.get_defined_type(&ctx.type_of_expr(e2)?);
                let w1 = self.compile_expr(e1, ctx)?;
                let w2 = self.compile_expr(e2, ctx)?;
                self.compile_inferred_binop(*op, t1, t2, w1, w2)
            }
            ast::Expr::UniOp(op, e) => {
                let t = ctx.get_defined_type(&ctx.type_of_expr(e)?);
                let w = self.compile_expr(e, ctx)?;
                self.compile_inferred_uniop(*op, t, w)
            }

            ast::Expr::Structure(_name, _vals) => {
                // Okay I'm going to put this off until I can handle struct layout
                // stuff.

                // First we compile all the member expressions
                /*
                for (_name, e) in vals.iter() {
                    self.compile_expr(e, ctx);
                }
                 */
                unimplemented!()
            }
            _ => unimplemented!(),
        };
        Ok(result_word)
    }

    /// Infers the proper instruction to emit based on the types given and the op,
    /// and outputs the appropriate instruction, returning its result address.
    /// So we can use `+` to add floats or vec4's and it works correctly.
    ///
    /// Panics if the types are not correct, such as adding a float and a bool.
    pub fn compile_inferred_binop(
        &mut self,
        op: ast::Op,
        t1: &verify::TypeDef,
        t2: &verify::TypeDef,
        e1: spirv::Word,
        e2: spirv::Word,
    ) -> spirv::Word {
        let f_word = self.get_type(&verify::TypeDef::F32);
        let b_word = self.get_type(&verify::TypeDef::Bool);
        match (op, t1, t2) {
            (ast::Op::Add, verify::TypeDef::F32, verify::TypeDef::F32) => {
                self.b.fadd(f_word, None, e1, e2).expect("???")
            }
            (ast::Op::Sub, verify::TypeDef::F32, verify::TypeDef::F32) => {
                self.b.fsub(f_word, None, e1, e2).expect("???")
            }
            (ast::Op::Mul, verify::TypeDef::F32, verify::TypeDef::F32) => {
                self.b.fmul(f_word, None, e1, e2).expect("???")
            }
            (ast::Op::Div, verify::TypeDef::F32, verify::TypeDef::F32) => {
                self.b.fdiv(f_word, None, e1, e2).expect("???")
            }

            (ast::Op::Gt, verify::TypeDef::F32, verify::TypeDef::F32) => {
                self.b.ford_greater_than(b_word, None, e1, e2).expect("???")
            }
            (ast::Op::Lt, verify::TypeDef::F32, verify::TypeDef::F32) => {
                self.b.ford_less_than(b_word, None, e1, e2).expect("???")
            }
            (ast::Op::Gte, verify::TypeDef::F32, verify::TypeDef::F32) => self
                .b
                .ford_greater_than_equal(b_word, None, e1, e2)
                .expect("???"),
            (ast::Op::Lte, verify::TypeDef::F32, verify::TypeDef::F32) => self
                .b
                .ford_less_than_equal(b_word, None, e1, e2)
                .expect("???"),
            (ast::Op::Eq, verify::TypeDef::F32, verify::TypeDef::F32) => {
                self.b.ford_equal(b_word, None, e1, e2).expect("???")
            }
            (ast::Op::Neq, verify::TypeDef::F32, verify::TypeDef::F32) => {
                self.b.ford_not_equal(b_word, None, e1, e2).expect("???")
            }

            (ast::Op::And, verify::TypeDef::Bool, verify::TypeDef::Bool) => {
                self.b.logical_and(b_word, None, e1, e2).expect("???")
            }
            (ast::Op::Or, verify::TypeDef::Bool, verify::TypeDef::Bool) => {
                self.b.logical_or(b_word, None, e1, e2).expect("???")
            }

            _ => {
                let msg = format!("Invalid binary op type: {:?} {:?} {:?}", op, t1, t2);
                panic!(msg)
            }
        }
    }

    /// Same as compile_inferred_binop but for unary operations
    pub fn compile_inferred_uniop(
        &mut self,
        op: ast::UOp,
        t: &verify::TypeDef,
        e: spirv::Word,
    ) -> spirv::Word {
        let f_word = self.get_type(&verify::TypeDef::F32);
        let b_word = self.get_type(&verify::TypeDef::Bool);
        match (op, t) {
            (ast::UOp::Negate, verify::TypeDef::F32) => {
                self.b.fnegate(f_word, None, e).expect("???")
            }
            (ast::UOp::Not, verify::TypeDef::Bool) => {
                self.b.logical_not(b_word, None, e).expect("???")
            }
            _ => {
                let msg = format!("Invalid unary op type: {:?} {:?}", op, t);
                panic!(msg)
            }
        }
    }

    pub fn compile_function(
        &mut self,
        ctx: &verify::VContext,
        def: &verify::FunctionDef,
    ) -> Result<(), crate::Error> {
        assert!(def.decl.body.len() > 0, "Empty function body!");
        let function_returns: spirv::Word =
            self.get_type(ctx.types.get(&def.decl.returns).unwrap());
        // Here we can assume all types exist.
        // TODO: We can make redundant function types here, maybe scrunch it down...
        // Not sure if SPIR-V can really have function pointers though so idk if they
        // can be first-class objects.
        // I think this is good now, we treat function types like all the other types.
        let ftype = self.add_type(&def.functiontype);

        let f_word = self.b.begin_function(
            function_returns,
            None,
            spirv::FunctionControl::DONT_INLINE
                | spirv::FunctionControl::CONST
                | spirv::FunctionControl::PURE,
            ftype,
        )?;

        // Bind the function name in the global symbol table.
        // This gets done before the body is evaluated so functions
        // can recurse.  Must be done before push_scope()!
        self.bind_var(&def.decl.name, f_word, ftype);

        self.push_scope();
        for p in def.decl.params.iter() {
            let param_typedef = ctx.get_defined_type(&p.typ);
            let param_typeword = self.get_type(param_typedef);
            let word = self.b.function_parameter(param_typeword)?;
            self.bind_var(&p.name, word, param_typeword);
        }

        self.b.begin_basic_block(None).unwrap();

        let last_expr_val: spirv::Word = def
            .decl
            .body
            .iter()
            .map(|e| self.compile_expr(e, ctx).expect("Could not compile expr!"))
            .last()
            .expect("Empty function body!  Try returning ()");
        self.b.ret_value(last_expr_val)?;
        self.b.end_function()?;
        self.pop_scope();

        // Add debug info
        self.b.name(f_word, &def.decl.name);

        Ok(())
    }

    /// Okay, so SPIR-V is a bit persnickity about entry points.
    /// An entry point MUST be a function that has no parameters
    /// or returns, and must describe the things it reads (and writes?)
    /// via links to particular global variables.
    ///
    /// I really don't feel like trying to handle that right now,
    /// as it will take a bit of squirrelly rewriting of basic function
    /// stuff to detect entry points and turn them into this special case
    /// form instead of just passing and returning values.  So for now
    /// we just stub out fake entry points.  Might be simplest to just
    /// make them call our actual entry point functions!
    /// Especially since entry points are defined to not be callable by
    /// anything else.
    ///
    /// TODO: Handle it!
    fn mongle_entry_points(&mut self, _ctx: &verify::VContext) -> Result<(), crate::Error> {
        let ftype = self.add_type(&verify::TypeDef::Function(
            vec![],
            Box::new(verify::TypeDef::Unit),
        ));
        let void = self.add_type(&verify::TypeDef::Unit);

        // Actually make our stub functions
        let v_word = self.b.begin_function(
            void,
            None,
            spirv::FunctionControl::DONT_INLINE | spirv::FunctionControl::CONST,
            ftype,
        )?;
        self.b.begin_basic_block(None)?;
        self.b.ret()?;
        self.b.end_function()?;

        let f_word = self.b.begin_function(
            void,
            None,
            spirv::FunctionControl::DONT_INLINE | spirv::FunctionControl::CONST,
            ftype,
        )?;
        self.b.begin_basic_block(None)?;
        self.b.ret()?;
        self.b.end_function()?;

        // We're allowed to reuse the same function as different
        // entry points, so I guess we'll just do that for now
        let vert_name = "_vertex_entry";
        let frag_name = "_fragment_entry";
        self.b.name(v_word, vert_name);
        self.b.name(f_word, frag_name);
        self.b
            .entry_point(spirv::ExecutionModel::Vertex, v_word, vert_name, []);
        self.b
            .entry_point(spirv::ExecutionModel::Fragment, f_word, frag_name, []);
        // TODO: Heck, what the heck to do with this??
        // Well, GLSL uses OriginUpperLeft, so let's do that for now.
        self.b
            .execution_mode(f_word, spirv::ExecutionMode::OriginUpperLeft, []);

        Ok(())

        /*
        // If this function is an entry point we declare it such.
        // TODO: This is a little jank but works.
        if def.decl.name == "vertex" {
            // Now if the function is an entry point we need to also add
            // global variables that form the inputs to it...
            self.b
                .entry_point(spirv::ExecutionModel::Vertex, f_word, "vertex", []);
        } else if def.decl.name == "fragment" {
            self.b
                .entry_point(spirv::ExecutionModel::Fragment, f_word, "vertex", []);
        }
        */
    }
}

pub fn compile(ctx: &verify::VContext) -> Result<CContext, crate::Error> {
    let mut cc = CContext::new();
    for (_name, def) in ctx.types.iter() {
        let _ = cc.add_type(def);
    }
    for (_name, def) in ctx.functions.iter() {
        cc.compile_function(ctx, def)?;
    }
    cc.mongle_entry_points(ctx)?;
    // TODO:
    // OpSource
    // OpName, at least for functions
    // OpMemberName for struct's
    // OpCapability
    Ok(cc)
}