use pyo3::Py;
use rspirv::{
    binary::Assemble,
    dr::{self, Operand},
    spirv,
};
use std::collections::HashMap;

use crate::pybytecode::{FunctionDef, Instruction, PossibleConsts, ShaderInput};

#[pyo3::pyfunction]
pub fn spv_gen(kernel: (u32, u32, u32), input: Py<ShaderInput>) -> Vec<u8> {
    let input = input.get();
    let main_name = input.main_function.as_str();

    let mut func_map: HashMap<String, &FunctionDef> = HashMap::new();
    for func in &input.functions {
        let f = func.get();
        func_map.insert(f.name.clone(), f);
    }

    let main_func = func_map.get(main_name).unwrap();

    let mut b = rspirv::dr::Builder::new();
    b.set_version(1, 3);
    b.capability(spirv::Capability::Shader);
    b.memory_model(spirv::AddressingModel::Logical, spirv::MemoryModel::GLSL450);

    let void = b.type_void();
    let voidf = b.type_function(void, vec![]);

    let int32 = b.type_int(32, 1);
    let float32 = b.type_float(32, None);

    let mut param_vars: HashMap<String, (u32, u32, u32, bool)> = HashMap::new();

    for (i, param) in main_func.inputs.iter().enumerate() {
        let p = param.get();
        let is_array = p.ty.starts_with("list[");
        let base_type = match p.ty.as_str() {
            "float" | "list[float]" => float32,
            "int" | "list[int]" => int32,
            _ => int32,
        };
        let member_type = if is_array {
            let runtime_array = b.type_runtime_array(base_type);
            b.decorate(
                runtime_array,
                spirv::Decoration::ArrayStride,
                vec![Operand::LiteralBit32(4)],
            );
            runtime_array
        } else {
            base_type
        };
        let struct_type = b.type_struct(vec![member_type]);
        b.decorate(struct_type, spirv::Decoration::Block, vec![]);
        b.member_decorate(
            struct_type,
            0,
            spirv::Decoration::Offset,
            vec![Operand::LiteralBit32(0)],
        );
        let ptr_storage = b.type_pointer(None, spirv::StorageClass::StorageBuffer, struct_type);
        let input_buf = b.variable(ptr_storage, None, spirv::StorageClass::StorageBuffer, None);
        b.decorate(
            input_buf,
            spirv::Decoration::DescriptorSet,
            vec![Operand::LiteralBit32(0)],
        );
        b.decorate(
            input_buf,
            spirv::Decoration::Binding,
            vec![Operand::LiteralBit32(i as u32)],
        );
        param_vars.insert(p.name.clone(), (base_type, ptr_storage, input_buf, is_array));
    }

    let mainf = b
        .begin_function(
            void,
            None,
            spirv::FunctionControl::DONT_INLINE | spirv::FunctionControl::CONST,
            voidf,
        )
        .unwrap();
    b.begin_block(None).unwrap();

    generate_body(
        &mut b,
        &main_func.instructions.iter().map(|i| i.get()).collect::<Vec<&Instruction>>(),
        &param_vars,
    );

    b.ret().unwrap();
    b.end_function().unwrap();
    b.entry_point(
        spirv::ExecutionModel::GLCompute,
        mainf,
        "main",
        vec![],
    );
    b.execution_mode(
        mainf,
        spirv::ExecutionMode::LocalSize,
        vec![kernel.0, kernel.1, kernel.2],
    );
    b.module()
        .assemble()
        .into_iter()
        .flat_map(|item| item.to_le_bytes())
        .collect()
}

fn make_const(b: &mut rspirv::dr::Builder, type_id: u32, value: u32) -> u32 {
    let id = b.id();
    let inst = dr::Instruction::new(
        spirv::Op::Constant,
        Some(type_id),
        Some(id),
        vec![Operand::LiteralBit32(value)],
    );
    b.module_mut().types_global_values.push(inst);
    id
}

fn generate_body(
    b: &mut rspirv::dr::Builder,
    instrs: &[&Instruction],
    param_vars: &HashMap<String, (u32, u32, u32, bool)>,
) {
    let int32 = b.type_int(32, 1);
    let float32 = b.type_float(32, None);

    let mut stack: Vec<(u32, u32, Option<(u32, u32)>)> = Vec::new();
    let mut locals: HashMap<String, (u32, u32)> = HashMap::new();
    let mut array_refs: HashMap<String, (u32, u32)> = HashMap::new();

    for instr in instrs {
        match instr {
            Instruction::LoadConst { val } => {
                let val = val.get();
                match val {
                    PossibleConsts::Int(v) => {
                        let cid = make_const(b, int32, *v as u32);
                        stack.push((cid, int32, None));
                    }
                    PossibleConsts::Float(v) => {
                        let cid = make_const(b, float32, v.to_bits());
                        stack.push((cid, float32, None));
                    }
                    PossibleConsts::Bool(v) => {
                        let cid = make_const(b, int32, if *v { 1 } else { 0 });
                        stack.push((cid, int32, None));
                    }
                    PossibleConsts::Null() => {}
                    PossibleConsts::Str(_) => {}
                }
            }
            Instruction::LoadFast { name } => {
                if let Some(&(vid, vtype)) = locals.get(name) {
                    let aref = array_refs.get(name).copied();
                    stack.push((vid, vtype, aref));
                } else if let Some(&(base_type, _ptr_storage, input_buf, is_array)) = param_vars.get(name) {
                    let zero_const = make_const(b, int32, 0);
                    let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
                    let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const]).unwrap();
                    let loaded = b.load(base_type, None, ptr, None, vec![]).unwrap();
                    locals.insert(name.clone(), (loaded, base_type));
                    if is_array {
                        array_refs.insert(name.clone(), (input_buf, base_type));
                    }
                    let aref = if is_array { Some((input_buf, base_type)) } else { None };
                    stack.push((loaded, base_type, aref));
                }
            }
            Instruction::LoadFastLoadFast { names } => {
                for name in [&names.0, &names.1] {
                    if let Some(&(vid, vtype)) = locals.get(name) {
                        let aref = array_refs.get(name).copied();
                        stack.push((vid, vtype, aref));
                    } else if let Some(&(base_type, _ptr_storage, input_buf, is_array)) = param_vars.get(name) {
                        let zero_const = make_const(b, int32, 0);
                        let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
                        let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const]).unwrap();
                        let loaded = b.load(base_type, None, ptr, None, vec![]).unwrap();
                        locals.insert(name.clone(), (loaded, base_type));
                        if is_array {
                            array_refs.insert(name.clone(), (input_buf, base_type));
                        }
                        let aref = if is_array { Some((input_buf, base_type)) } else { None };
                        stack.push((loaded, base_type, aref));
                    }
                }
            }
            Instruction::StoreFast { name } => {
                if let Some((vid, vtype, _aref)) = stack.pop() {
                    locals.insert(name.clone(), (vid, vtype));
                    if let Some(&(base_type, _ptr_storage, input_buf, _is_array)) = param_vars.get(name) {
                        let zero_const = make_const(b, int32, 0);
                        let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
                        let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const]).unwrap();
                        b.store(ptr, vid, None, vec![]).unwrap();
                    }
                }
            }
            Instruction::BinOp { op } => {
                if let (Some((rhs, rhs_type, _)), Some((lhs, lhs_type, _))) = (stack.pop(), stack.pop()) {
                    let (res_type, res) = if rhs_type == float32 || lhs_type == float32 {
                        match op.as_str() {
                            "+" | "+=" => (float32, b.f_add(float32, None, lhs, rhs).unwrap()),
                            "-" | "-=" => (float32, b.f_sub(float32, None, lhs, rhs).unwrap()),
                            "*" | "*=" => (float32, b.f_mul(float32, None, lhs, rhs).unwrap()),
                            "/" | "//" | "/=" => (float32, b.f_div(float32, None, lhs, rhs).unwrap()),
                            "%" | "%=" => (float32, b.f_mod(float32, None, lhs, rhs).unwrap()),
                            _ => (float32, lhs),
                        }
                    } else {
                        match op.as_str() {
                            "+" | "+=" => (int32, b.i_add(int32, None, lhs, rhs).unwrap()),
                            "-" | "-=" => (int32, b.i_sub(int32, None, lhs, rhs).unwrap()),
                            "*" | "*=" => (int32, b.i_mul(int32, None, lhs, rhs).unwrap()),
                            "/" | "//" | "/=" => (int32, b.s_div(int32, None, lhs, rhs).unwrap()),
                            "%" | "%=" => (int32, b.s_rem(int32, None, lhs, rhs).unwrap()),
                            "**" => (int32, lhs),
                            _ => (int32, lhs),
                        }
                    };
                    stack.push((res, res_type, None));
                }
            }
            Instruction::UnaryOp { op } => {
                if let Some((val, vtype, _)) = stack.pop() {
                    let (res_type, res) = if vtype == float32 {
                        match op.as_str() {
                            "-" => (float32, b.f_negate(float32, None, val).unwrap()),
                            "!" => {
                                let bool_type = b.type_bool();
                                let zero = make_const(b, float32, 0.0f32.to_bits());
                                (bool_type, b.f_ord_equal(bool_type, None, val, zero).unwrap())
                            }
                            _ => (vtype, val),
                        }
                    } else {
                        match op.as_str() {
                            "-" => (int32, b.s_negate(int32, None, val).unwrap()),
                            "!" => {
                                let bool_type = b.type_bool();
                                let zero = make_const(b, int32, 0);
                                (bool_type, b.i_equal(bool_type, None, val, zero).unwrap())
                            }
                            _ => (vtype, val),
                        }
                    };
                    stack.push((res, res_type, None));
                }
            }
            Instruction::CompareOp { op } => {
                if let (Some((rhs, rhs_type, _)), Some((lhs, lhs_type, _))) = (stack.pop(), stack.pop()) {
                    let bool_type = b.type_bool();
                    let (res, res_type) = if rhs_type == float32 || lhs_type == float32 {
                        match op.as_str() {
                            "== " | "==" => (b.f_ord_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            "!=" => (b.f_ord_not_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            "<" => (b.f_ord_less_than(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            "<=" => (b.f_ord_less_than_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            ">" => (b.f_ord_greater_than(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            ">=" => (b.f_ord_greater_than_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            _ => (b.f_ord_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                        }
                    } else {
                        match op.as_str() {
                            "== " | "==" => (b.i_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            "!=" => (b.i_not_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            "<" => (b.s_less_than(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            "<=" => (b.s_less_than_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            ">" => (b.s_greater_than(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            ">=" => (b.s_greater_than_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                            _ => (b.i_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
                        }
                    };
                    stack.push((res, res_type, None));
                }
            }
            Instruction::PopTop {} => {
                stack.pop();
            }
            Instruction::ReturnValue {} => {
                stack.pop();
            }
            Instruction::LoadGlobal { .. } => {}
            Instruction::BinarySubscr {} => {
                if let (Some((idx, _, _)), Some((_, base_type, Some((input_buf, _))))) = (stack.pop(), stack.pop()) {
                    let zero_const = make_const(b, int32, 0);
                    let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
                    let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const, idx]).unwrap();
                    let loaded = b.load(base_type, None, ptr, None, vec![]).unwrap();
                    stack.push((loaded, base_type, None));
                }
            }
            Instruction::StoreSubscr {} => {
                if let (Some((idx, _, _)), Some((_, _, Some((input_buf, base_type)))), Some((val, _vtype, _))) = (stack.pop(), stack.pop(), stack.pop()) {
                    let zero_const = make_const(b, int32, 0);
                    let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
                    let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const, idx]).unwrap();
                    b.store(ptr, val, None, vec![]).unwrap();
                }
            }
            Instruction::Copy { n } => {
                let n = *n as usize;
                let len = stack.len();
                let idx = len - n;
                stack.push(stack[idx].clone());
            }
            Instruction::Swap { n } => {
                let n = *n as usize;
                let len = stack.len();
                stack.swap(len - 1, len - n);
            }
            Instruction::Call {} => {}
        }
    }
}