use heck::{ToShoutySnakeCase, ToSnakeCase, ToUpperCamelCase};
use proc_macro2::{Ident, Span, TokenStream};
use quote::{format_ident, quote};
use std::collections::HashMap;

use crate::spec::{Category, CategoryData, Grammar, InstructionData, Kind, Operand, Quantifier};

pub fn generate_rust_grammar(grammar: Grammar) -> TokenStream {
    let mut out = TokenStream::new();

    let mut kinds: Vec<_> = grammar.categories.iter().collect();
    kinds.sort_by_key(|((_, k), _)| format!("{k:?}"));

    for ((category, kind), data) in &kinds {
        match category {
            Category::BitEnum   => out.extend(gen_bit_enum(kind, data)),
            Category::ValueEnum => out.extend(gen_value_enum(kind, data)),
            _                   => {}
        }
    }

    let mut ops: Vec<_> = grammar.instructions.iter().collect();
    ops.sort_by_key(|(opcode, _)| *opcode);

    let opcode_variants = ops.iter().map(|(val, instr)| {
        let variant = opcode_variant_ident(&instr.opname);
        let v       = *val;
        quote! { #variant = #v, }
    });
    out.extend(quote! {
        #[repr(u32)]
        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
        pub enum Opcode {
            #(#opcode_variants)*
        }
    });

    let mut op_tokens = TokenStream::new();
    for (_, instr) in &ops {
        op_tokens.extend(gen_op(instr, &grammar.categories));
    }

    out.extend(quote! {
        pub mod ops {
            use super::*;
            #op_tokens
        }
    });

    out
}

// ── BitEnum ───────────────────────────────────────────────────────────────────

fn gen_bit_enum(kind: &Kind, data: &CategoryData) -> TokenStream {
    let name = kind_ident(kind);

    let mut enumerants: Vec<_> = data.enumerants.iter().collect();
    enumerants.sort_by_key(|(_, e)| e.value);

    let flags = enumerants.iter()
        .filter(|(_, e)| e.value != 0)
        .map(|(enumerant, e)| {
            let const_name = safe_ident(enumerant.to_shouty_snake_case());
            let val        = e.value;
            let aliases    = e.meta.aliases.iter().map(|a| quote!(#[doc(alias = #a)]));
            quote! { #(#aliases)* const #const_name = #val; }
        });

    let mut out = quote! {
        bitflags::bitflags! {
            #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
            #[repr(transparent)]
            pub struct #name: u32 {
                #(#flags)*
            }
        }
        impl Spv for #name {
            fn encode(&self, out: &mut Vec<u32>) { out.push(self.bits()); }
        }
    };

    let param_fields: Vec<_> = enumerants.iter()
        .filter(|(_, e)| e.value != 0 && !e.parameters.is_empty())
        .map(|(enumerant, e)| {
            let fname = safe_ident(enumerant.to_snake_case());
            let ftype = params_type(&e.parameters);
            let def   = quote! { pub #fname: Option<#ftype> };
            (fname, def)
        })
        .collect();

    if !param_fields.is_empty() {
        let params_name   = format_ident!("{}Params", format!("{kind:?}"));
        let field_defs    = param_fields.iter().map(|(_, def)| def);
        let field_encodes = param_fields.iter().map(|(name, _)| quote! { self.#name.encode(out); });
        out.extend(quote! {
            #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
            pub struct #params_name {
                pub flags: #name,
                #(#field_defs,)*
            }
            impl Spv for #params_name {
                fn encode(&self, out: &mut Vec<u32>) {
                    self.flags.encode(out);
                    #(#field_encodes)*
                }
            }
        });
    }

    out
}

// ── ValueEnum ─────────────────────────────────────────────────────────────────

fn gen_value_enum(kind: &Kind, data: &CategoryData) -> TokenStream {
    let name = kind_ident(kind);

    let mut enumerants: Vec<_> = data.enumerants.iter().collect();
    enumerants.sort_by_key(|(_, e)| e.value);

    let variants = enumerants.iter().map(|(enumerant, e)| {
        let vname   = safe_ident(enumerant.to_upper_camel_case());
        let aliases = e.meta.aliases.iter().map(|a| quote!(#[doc(alias = #a)]));
        let body    = match e.parameters.as_slice() {
            []  => quote! { #vname },
            [p] => { let t = kind_to_type(&p.kind); quote! { #vname(#t) } }
            ps  => { let ts: Vec<_> = ps.iter().map(|p| kind_to_type(&p.kind)).collect();
                     quote! { #vname(#(#ts),*) } }
        };
        quote! { #(#aliases)* #body, }
    });

    let as_u32_arms = enumerants.iter().map(|(enumerant, e)| {
        let vname = safe_ident(enumerant.to_upper_camel_case());
        let val   = e.value;
        if e.parameters.is_empty() {
            quote! { Self::#vname      => #val, }
        } else {
            quote! { Self::#vname(..) => #val, }
        }
    });

    let spv_arms = enumerants.iter().map(|(enumerant, e)| {
        let vname = safe_ident(enumerant.to_upper_camel_case());
        let val   = e.value;
        match e.parameters.as_slice() {
            [] => quote! { Self::#vname => { out.push(#val); } },
            [_] => quote! { Self::#vname(ref v) => { out.push(#val); v.encode(out); } },
            ps  => {
                let bindings: Vec<_> = (0..ps.len()).map(|i| format_ident!("v{i}")).collect();
                let encodes = bindings.iter().map(|b| quote! { #b.encode(out); });
                quote! { Self::#vname(#(ref #bindings),*) => { out.push(#val); #(#encodes)* } }
            }
        }
    });

    quote! {
        #[derive(Clone, Debug, PartialEq, Eq, Hash)]
        pub enum #name {
            #(#variants)*
        }
        impl #name {
            pub fn as_u32(&self) -> u32 {
                match self { #(#as_u32_arms)* }
            }
        }
        impl Spv for #name {
            fn encode(&self, out: &mut Vec<u32>) {
                match self { #(#spv_arms)* }
            }
        }
    }
}

// ── Op structs ────────────────────────────────────────────────────────────────

fn gen_op(instr: &InstructionData, categories: &HashMap<(Category, Kind), CategoryData>) -> TokenStream {
    let name    = safe_ident(instr.opname.clone());
    let variant = opcode_variant_ident(&instr.opname);

    let has_result_type = instr.operands.iter().any(|o| o.kind == Kind::IdResultType);
    let has_result_id   = instr.operands.iter().any(|o| o.kind == Kind::IdResult);

    let mut seen         = HashMap::new();
    let mut fields       = Vec::new();
    let mut field_idents = Vec::<TokenStream>::new();

    if has_result_type {
        fields.push(quote! { pub result_type: Id });
        field_idents.push(quote! { result_type });
    }
    if has_result_id {
        fields.push(quote! { pub result_id: Id });
        field_idents.push(quote! { result_id });
    }

    for (i, op) in instr.operands.iter()
        .filter(|o| o.kind != Kind::IdResult && o.kind != Kind::IdResultType)
        .enumerate()
    {
        let fname = dedup(field_name(op, i), &mut seen);
        let ftype = if bit_enum_has_params(&op.kind, categories) {
            let pname = format_ident!("{}Params", format!("{:?}", op.kind));
            wrap_q(quote!(#pname), &op.quantifier)
        } else {
            wrap_q(kind_to_type(&op.kind), &op.quantifier)
        };
        field_idents.push(quote! { #fname });
        fields.push(quote! { pub #fname: #ftype });
    }

    let encodes = field_idents.iter().map(|f| quote! { self.#f.encode(out); });

    let result_id_expr   = if has_result_id   { quote!(Some(self.result_id))   } else { quote!(None) };
    let result_type_expr = if has_result_type { quote!(Some(self.result_type)) } else { quote!(None) };

    quote! {
        #[derive(Clone, Debug, PartialEq, Eq, Hash)]
        pub struct #name {
            #(#fields,)*
        }
        impl #name {
            pub const OPCODE: Opcode = Opcode::#variant;
        }
        impl Spv for #name {
            fn encode(&self, out: &mut Vec<u32>) {
                let start = out.len();
                out.push(0u32);
                #(#encodes)*
                let wc = (out.len() - start) as u32;
                out[start] = Opcode::#variant as u32 | (wc << 16);
            }
        }
        impl Instruction for #name {
            fn opcode(&self)      -> Opcode     { Opcode::#variant }
            fn result_id(&self)   -> Option<Id> { #result_id_expr }
            fn result_type(&self) -> Option<Id> { #result_type_expr }
        }
        impl From<#name> for Opcode {
            fn from(_: #name) -> Self { Opcode::#variant }
        }
    }
}

// ── helpers ───────────────────────────────────────────────────────────────────

fn kind_to_type(kind: &Kind) -> TokenStream {
    match kind {
        Kind::IdResultType | Kind::IdResult | Kind::IdRef
        | Kind::IdScope    | Kind::IdMemorySemantics           => quote!(Id),
        Kind::LiteralInteger | Kind::LiteralContextDependentNumber
        | Kind::LiteralExtInstInteger | Kind::LiteralSpecConstantOpInteger
        | Kind::LiteralFloat                                   => quote!(u32),
        Kind::LiteralString                                    => quote!(::std::string::String),
        Kind::PairLiteralIntegerIdRef                          => quote!((u32, Id)),
        Kind::PairIdRefLiteralInteger                          => quote!((Id, u32)),
        Kind::PairIdRefIdRef                                   => quote!((Id, Id)),
        _                                                      => { let i = kind_ident(kind); quote!(#i) }
    }
}

fn wrap_q(ty: TokenStream, q: &Quantifier) -> TokenStream {
    match q {
        Quantifier::Required   => ty,
        Quantifier::Optional   => quote!(Option<#ty>),
        Quantifier::ZeroOrMore => quote!(Vec<#ty>),
    }
}

fn params_type(params: &[Operand]) -> TokenStream {
    let ts: Vec<_> = params.iter().map(|p| kind_to_type(&p.kind)).collect();
    if ts.len() == 1 { ts.into_iter().next().unwrap() } else { quote!((#(#ts),*)) }
}

fn field_name(op: &Operand, _i: usize) -> String {
    let raw = if op.name.is_empty() { format!("{:?}", op.kind) } else { op.name.clone() };
    let raw = raw.split(',').next().unwrap().trim().to_string();
    let raw = raw.trim_end_matches(|c: char| c.is_ascii_digit() || c == ' ').to_string();
    let s   = raw.replace('\'', "").to_snake_case();
    match s.as_str() {
        "type" => "ty".into(),
        "ref"  => "ref_".into(),
        "use"  => "use_".into(),
        _      => s,
    }
}

fn dedup(name: String, seen: &mut HashMap<String, usize>) -> Ident {
    let n = seen.entry(name.clone()).or_insert(0);
    let s = if *n == 0 { name } else { format!("{name}_{n}") };
    *n += 1;
    safe_ident(s)
}

fn bit_enum_has_params(kind: &Kind, categories: &HashMap<(Category, Kind), CategoryData>) -> bool {
    categories.iter()
        .find(|((cat, k), _)| cat == &Category::BitEnum && k == kind)
        .map(|(_, d)| d.enumerants.values().any(|e| !e.parameters.is_empty()))
        .unwrap_or(false)
}

fn opcode_variant_ident(opname: &str) -> Ident {
    safe_ident(opname.strip_prefix("Op").unwrap_or(opname).to_string())
}

fn kind_ident(kind: &Kind) -> Ident {
    Ident::new(&format!("{kind:?}"), Span::call_site())
}

fn safe_ident(s: String) -> Ident {
    if s.starts_with(|c: char| c.is_ascii_digit()) {
        format_ident!("_{}", s)
    } else {
        format_ident!("{}", s)
    }
}
