#[cfg(test)]
mod tests;

use crate::compiler::{
    AddressingModel, BuiltinDecoration, Decoration, ExecutionMode, ExecutionModel, Instruction,
    MemoryModel, Module, StorageClass,
};
use std::fmt;

use crate::parser::Ast;

#[derive(Debug, PartialEq, Clone)]
pub enum Type {
    Pointer(Box<Type>, StorageClass),
    Vector(Box<Type>, u32),
    Float(u32),
    Int(u32),
    Unsigned(u32),
    Void,
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Type::Void => write!(f, "<>"),
            Type::Pointer(typ, storage_class) => match storage_class {
                StorageClass::Input => write!(f, "*{}i", typ),
                StorageClass::Output => write!(f, "*{}o", typ),
                StorageClass::Undefined => panic!("Bound a non-declared variable"),
            },
            Type::Vector(typ, size) => write!(f, "v{}{}", size, typ),
            Type::Float(size) => write!(f, "f{}", size),
            Type::Int(size) => write!(f, "s{}", size),
            Type::Unsigned(size) => write!(f, "u{}", size),
        }
    }
}

pub fn parse_type(typ: &String) -> Type {
    // pointers have the format *<t>i or *<t>o, for the storage class
    // floats have the format f<size>, so f32, f64...
    // ints have the format s<size>, so s32, s64...
    // unsigned ints have the format u<size>, so u32, u64...
    // So, *v4f32i is a pointer to a vector of 4 floats, with the storage class Input.
    // *v4f32o is a pointer to a vector of 4 floats, with the storage class Output.

    let mut chars = typ.chars();
    match chars.next() {
        Some('<') => {
            assert_eq!(typ.as_str(), "<>");
            Type::Void
        }
        Some('*') => {
            let mut typ = chars.collect::<String>();
            let storage_class = match typ.pop() {
                Some('i') => StorageClass::Input,
                Some('o') => StorageClass::Output,
                _ => panic!("Invalid storage class"),
            };
            Type::Pointer(Box::new(parse_type(&typ)), storage_class)
        }
        Some('v') => {
            let size = chars.next().map(|s| s.to_digit(10)).flatten().unwrap();
            let typ = chars.collect::<String>();
            Type::Vector(Box::new(parse_type(&typ)), size)
        }
        Some('f') => {
            let size = chars.collect::<String>().parse().unwrap();
            Type::Float(size)
        }
        Some('s') => {
            let size = chars.collect::<String>().parse().unwrap();
            Type::Int(size)
        }
        Some('u') => {
            let size = chars.collect::<String>().parse().unwrap();
            Type::Unsigned(size)
        }
        _ => panic!("Invalid type"),
    }
}

fn emit_type(typ: &Type, ops: &mut Vec<(String, String)>) {
    match &typ {
        Type::Void => {
            ops.push(("%void".to_string(), "OpTypeVoid".to_string()));
        }
        Type::Unsigned(size) => {
            ops.push((
                format!("%u{}", size).to_string(),
                format!("OpTypeInt {}", size),
            ));
        }
        Type::Int(size) => {
            ops.push((
                format!("%s{}", size).to_string(),
                format!("OpTypeInt {}", size),
            ));
        }
        Type::Float(size) => {
            ops.push((
                format!("%f{}", size).to_string(),
                format!("OpTypeFloat {}", size),
            ));
        }
        Type::Vector(in_typ, size) => {
            emit_type(in_typ, ops);
            ops.push((
                fix_name(&typ.to_string()),
                format!("OpTypeVector {} {}", fix_name(&in_typ.to_string()), size),
            ))
        }
        Type::Pointer(in_typ, storage_class) => {
            emit_type(in_typ, ops);
            let typ_id = fix_name(&typ.to_string());
            let storage_class = match storage_class {
                StorageClass::Input => "Input",
                StorageClass::Output => "Output",
                StorageClass::Undefined => panic!("Bound a non-declared variable"),
            };
            ops.push((
                typ_id,
                format!(
                    "OpTypePointer {} {}",
                    storage_class,
                    fix_name(&in_typ.to_string())
                ),
            ))
        }
    }
}

fn fix_name(name: &str) -> String {
    format!(
        "%{}",
        name.replace("-", "_")
            .replace("*", "p")
            .replace("<>", "void")
    )
}

fn has_id<T>(name: &str, ops: &Vec<(Option<String>, T)>) -> bool {
    ops.iter()
        .find(|op| op.0.as_ref().map(|s| *s == name).unwrap_or(false))
        .is_some()
}

pub fn spirv_meta(module: &mut Module) -> Vec<(Option<String>, Vec<String>)> {
    let mut ops: Vec<(Option<String>, Vec<String>)> = Vec::new();

    let capabilities: Vec<String> = module
        .capabilities
        .iter()
        .map(|c| format!("{:?}", c))
        .collect();
    for cap in capabilities {
        ops.push((None, vec!["OpCapability".to_string(), cap]));
    }

    let memory_model_address = match module.memory_model.addressing_model {
        AddressingModel::Logical => "Logical",
        AddressingModel::Physical32 => "Physical32",
        AddressingModel::Physical64 => "Physical64",
        AddressingModel::PhysicalStorageBuffer64 => "PhysicalStorageBuffer64",
    };
    let memory_model_model = match module.memory_model.memory_model {
        MemoryModel::Simple => "Simple",
        MemoryModel::GLSL450 => "GLSL450",
        MemoryModel::OpenCL => "OpenCL",
        _ => todo!(),
    };
    ops.push((
        None,
        vec![
            "OpMemoryModel".to_string(),
            memory_model_address.to_string(),
            memory_model_model.to_string(),
        ],
    ));

    for entry in module.entry_points.clone() {
        let exec_model = match entry.execution_model {
            ExecutionModel::Fragment => "Fragment",
            ExecutionModel::Vertex => "Vertex",
        };
        let name = entry.name;
        let interface: Vec<String> = entry
            .interface
            .iter()
            .map(|i| fix_name(&i.to_string()))
            .collect();
        let exec_mode = match entry.execution_mode {
            ExecutionMode::OriginUpperLeft => "OriginUpperLeft",
        };
        ops.push((
            None,
            vec![
                "OpEntryPoint".to_string(),
                exec_model.to_string(),
                fix_name(&name),
                format!("\"{}\"", name),
                interface.join(" "),
            ],
        ));
        ops.push((
            None,
            vec![
                "OpExecutionMode".to_string(),
                fix_name(&name),
                exec_mode.to_string(),
            ],
        ));
    }

    for global in module.globals.clone() {
        let name = fix_name(&global.name);
        let typ = global.typ;
        let storage_class = match global.storage_class {
            StorageClass::Input => "Input",
            StorageClass::Output => "Output",
            StorageClass::Undefined => panic!("Bound a non-declared variable"),
        };
        let mut type_ops = Vec::new();
        emit_type(&parse_type(&typ), &mut type_ops);
        for op in type_ops {
            if has_id(&op.0, &ops) {
                continue;
            }
            ops.push((Some(op.0), vec![op.1]));
        }
        ops.push((
            Some(name.clone()),
            vec![
                "OpVariable".to_string(),
                fix_name(&typ),
                storage_class.to_string(),
            ],
        ));
        for dec in global.decorations {
            // Decorations have the format Location 0, or Builtin FragCoord
            let dec = match dec {
                Decoration::Location(loc) => format!("Location {}", loc),
                Decoration::BuiltIn(builtin) => {
                    let builtin = match builtin {
                        BuiltinDecoration::FragCoord => "FragCoord",
                    };
                    format!("BuiltIn {}", builtin)
                }
            };
            ops.push((None, vec!["OpDecorate".to_string(), name.clone(), dec]));
        }
    }

    for fun in module.functions.clone() {
        let name = fix_name(&format!("l_{}", fun.name));
        let return_type = fix_name(&fun.return_type);
        let mut type_ops = Vec::new();
        emit_type(&parse_type(&fun.return_type), &mut type_ops);
        for op in type_ops {
            if has_id(&op.0, &ops) {
                continue;
            }
            ops.push((Some(op.0), vec![op.1]));
        }
        // Push OpFunctionType
        ops.push((Some(name), vec!["OpTypeFunction".to_string(), return_type]));
    }

    // for op in ops {
    //     if let Some(op0) = op.0 {
    //         write!(spirv_asm, "{} = ", op0).unwrap();
    //     }
    //     for arg in op.1 {
    //         write!(spirv_asm, "{} ", arg).unwrap();
    //     }
    //     writeln!(spirv_asm).unwrap();
    // }
    ops
}

enum Number {
    Int(i32),
    Float(f32),
    NotANumber,
}

fn match_number(s: &str) -> Number {
    // floats have to be in the format of [0-9]+\.[0-9]*
    // integers have to be in the format of [0-9]+
    let mut chars = s.chars();
    let mut has_dot = false;
    while let Some(c) = chars.next() {
        if c == '.' {
            if has_dot {
                // only one dot allowed.
                return Number::NotANumber;
            }
            has_dot = true;
        } else if !c.is_digit(10) {
            // not a number;
            // has a character that is not a digit or a dot.
            return Number::NotANumber;
        }
    }
    if has_dot {
        // we have checked that the whole thing is numbers and dots
        // and that there is precisely one dot
        // that matches the regex.
        Number::Float(s.parse().unwrap())
    } else {
        // this is an integer, since no dots.
        Number::Int(s.parse().unwrap())
    }
}

fn compile_biop(
    op: &str,
    lst: &mut Vec<Ast>,
    vars: &mut Vec<(String, String)>,
    constants: &mut Vec<(String, String)>,
    types: &mut Vec<String>,
    counter: &mut i32,
    stack: &mut Vec<String>,
    out: &mut Vec<(Option<String>, String)>,
) {
    assert!(lst.len() == 2);
    let rhs = lst.pop().unwrap();
    let lhs = lst.pop().unwrap();
    compile_ast_ssa(lhs, vars, constants, types, counter, stack, out);
    compile_ast_ssa(rhs, vars, constants, types, counter, stack, out);
    let rhs_id = stack.pop().unwrap();
    let lhs_id = stack.pop().unwrap();
    let id = String::from(counter.to_string());
    *counter += 1;
    out.push((
        Some(id.clone()),
        format!(
            "{} {} {} {}",
            op,
            fix_name(&String::from("f32")),
            fix_name(&lhs_id),
            fix_name(&rhs_id),
        ),
    ));
    stack.push(id);
}

pub fn compile_ast_ssa(
    ast: Ast,
    vars: &mut Vec<(String, String)>,
    constants: &mut Vec<(String, String)>,
    types: &mut Vec<String>,
    counter: &mut i32,
    stack: &mut Vec<String>,
    out: &mut Vec<(Option<String>, String)>,
) {
    match ast.clone().list() {
        Some(mut lst) => {
            assert!(!lst.is_empty());
            let fun = lst.remove(0);
            assert!(true); // no safe remove, thanks Rust
            let fun_name = fun.symbol();
            assert!(fun_name.is_some());
            let fun_name = fun_name.unwrap();
            match fun_name.as_str() {
                "store-ptr" => {
                    assert!(lst.len() == 2);
                    let ptr = lst.pop().unwrap();
                    let val = lst.pop().unwrap();
                    compile_ast_ssa(ptr, vars, constants, types, counter, stack, out);
                    compile_ast_ssa(val, vars, constants, types, counter, stack, out);
                    let val_id = stack.pop().unwrap();
                    let ptr_id = stack.pop().unwrap();
                    out.push((
                        None,
                        format!("OpStore {} {}", fix_name(&val_id), fix_name(&ptr_id)),
                    ));
                }
                "/" => {
                    compile_biop(
                        "OpFDiv", &mut lst, vars, constants, types, counter, stack, out,
                    );
                }
                "*" => {
                    compile_biop(
                        "OpFMul", &mut lst, vars, constants, types, counter, stack, out,
                    );
                }
                "+" => {
                    compile_biop(
                        "OpFAdd", &mut lst, vars, constants, types, counter, stack, out,
                    );
                }
                "-" => {
                    compile_biop(
                        "OpFSub", &mut lst, vars, constants, types, counter, stack, out,
                    );
                }
                s => {
                    panic!(
                        "Unknown function: {} with params {:#?} in context:\n{:#?}",
                        s, lst, ast
                    );
                }
            }
        }
        None => {
            let sym = ast.symbol().unwrap();
            match match_number(&sym) {
                Number::Int(i) => {
                    let key = format!("i32_{}", i);
                    let mut contains = false;
                    for c in constants.iter() {
                        if c.0 == key {
                            contains = true;
                        }
                    }
                    if !contains {
                        constants.push((key.clone(), format!("OpConstant %i32 {}", i.to_string())));
                    }
                    stack.push(key);
                }
                Number::Float(f) => {
                    let key = format!("f32_{}", f);
                    let mut contains = false;
                    for c in constants.iter() {
                        if c.0 == key {
                            contains = true;
                        }
                    }
                    if !contains {
                        constants.push((key.clone(), format!("OpConstant %f32 {}", f.to_string())));
                    }
                    for t in types.iter() {
                        if t == "f32" {
                            contains = true;
                        }
                    }
                    if !contains {
                        types.push("f32".to_string());
                    }
                    stack.push(key);
                }
                Number::NotANumber => {
                    for v in vars.iter() {
                        if v.0 == sym {
                            stack.push(v.0.clone());
                            return;
                        }
                    }
                    panic!("Unknown variable or constant: {}", sym);
                }
            }
        }
    }
}

pub fn compile_fun_ssa(module: &mut Module, ops: &Vec<(Option<String>, String)>) {
    let mut label_counter = Box::new(0);
    for fun in module.functions.iter_mut() {
        assert!(fun.ast.is_some());
        let ast = fun.ast.as_mut().unwrap();
        let block = ast.clone().list().unwrap().get(0).unwrap().clone();
        let mut vars = vec![];
        let mut constants = vec![];
        let mut types = vec![];
        let mut counter = Box::new(0);
        let mut stack = vec![];
        let mut out_op = vec![];
        for v in &module.globals {
            vars.push((v.name.clone(), v.typ.clone()));
        }
        compile_ast_ssa(
            block,
            &mut vars,
            &mut constants,
            &mut types,
            &mut counter,
            &mut stack,
            &mut out_op,
        );
        let mut out_pre = vec![];
        for t in &types {
            let typ = parse_type(t);
            if has_id(&fix_name(&typ.to_string()), ops) {
                continue;
            }
            if has_id(&fix_name(&typ.to_string()), &out_pre) {
                continue;
            }
            let mut type_ops = vec![];
            emit_type(&typ, &mut type_ops);
            for type_op in type_ops {
                out_pre.push((Some(type_op.0), type_op.1));
            }
        }
        for c in &constants {
            out_pre.push((Some(fix_name(&c.0.clone())), c.1.clone()));
        }

        // TODO non-void type
        out_pre.push((
            Some(fix_name(&fun.name)),
            format!("OpFunction %void None %l_{}", fun.name),
        ));
        out_pre.push((
            Some(format!("%n_{}", label_counter.to_string())),
            "OpLabel".to_string(),
        ));
        *label_counter += 1;
        let mut out_ops = out_pre.clone();
        for op in out_op {
            if op.0.is_some() {
                out_ops.push((Some(fix_name(&op.0.unwrap())), op.1));
            } else {
                out_ops.push((None, op.1));
            }
        }
        for op in out_ops {
            let split: Vec<String> = op.1.split(" ").map(|s| s.to_string()).collect();
            let op_name: String = (&split[0]).clone();
            let op_args: Vec<String> = split[1..].iter().map(|s| s.clone()).collect();
            let op_id = op.0.clone();
            let ins: Instruction = Instruction {
                result_id: op_id,
                op: op_name,
                operands: op_args,
            };
            fun.body.as_mut().unwrap().push(ins);
        }
    }
}