use serde::{
    de::{self, MapAccess, Visitor},
    Deserialize, Deserializer,
};
use std::collections::HashMap;
use std::fmt;

const SPIRV_CORE_GRAMMAR: &str = include_str!("spec/spirv.core.grammar.json");

pub struct Mapping<'a> {
    pub category_spec: &'a str,
    pub module_name:   &'a str,
    pub enum_name:     &'a str,
}

macro_rules! map {
    ($cat:literal => $enum:literal in $mod:literal) => {
        Mapping {
            category_spec: $cat,
            module_name: $mod,
            enum_name: $enum,
        }
    };
}

const GRAMMAR_MAPPING: [Mapping<'static>; 17] = [
    map!("Miscellaneous"          => "Misc"        in "misc"),
    map!("Debug"                  => "Debug"       in "debug"),
    map!("Type-Declaration"       => "Type"        in "ty"),
    map!("Constant-Creation"      => "Const"       in "co"),
    map!("Memory"                 => "Mem"         in "mem"),
    map!("Function"               => "Func"        in "func"),
    map!("Image"                  => "Image"       in "image"),
    map!("Conversion"             => "Convert"     in "convert"),
    map!("Composite"              => "Composite"   in "composite"),
    map!("Arithmetic"             => "Math"        in "math"),
    map!("Derivative"             => "Math"        in "math"),
    map!("Bit"                    => "Bit"         in "bit"),
    map!("Relational_and_Logical" => "Relation"    in "relation"),
    map!("Control-Flow"           => "ControlFlow" in "control_flow"),
    map!("Atomic"                 => "Atomic"      in "atomic"),
    map!("Barrier"                => "Barrier"     in "barrier"),
    map!("Group"                  => "Group"       in "group"),
];

// ── Version ───────────────────────────────────────────────────────────────────

#[derive(Debug, Default, PartialEq, Eq)]
pub struct Version(pub u8, pub u8, pub u8);

struct VersionVisitor;

impl<'de> Visitor<'de> for VersionVisitor {
    type Value = Version;

    fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, r#"a version string like "1.0""#)
    }

    fn visit_str<E: de::Error>(self, s: &str) -> Result<Version, E> {
        let (maj, min) = s
            .split_once('.')
            .ok_or_else(|| de::Error::custom(format!("invalid version: {s}")))?;
        Ok(Version(
            maj.parse().map_err(de::Error::custom)?,
            min.parse().map_err(de::Error::custom)?,
            0,
        ))
    }
}

impl<'de> Deserialize<'de> for Version {
    fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        d.deserialize_str(VersionVisitor)
    }
}

fn de_opt_version<'de, D: Deserializer<'de>>(d: D) -> Result<Option<Version>, D::Error> {
    let s = String::deserialize(d)?;
    if s == "None" {
        return Ok(None);
    }
    VersionVisitor.visit_str::<D::Error>(&s).map(Some)
}

// ── EnumerantValue ────────────────────────────────────────────────────────────

#[derive(Deserialize)]
#[serde(untagged)]
enum RawValue {
    Str(String),
    Int(u32),
}

fn de_value<'de, D: Deserializer<'de>>(d: D) -> Result<u32, D::Error> {
    match RawValue::deserialize(d)? {
        RawValue::Int(n) => Ok(n),
        RawValue::Str(s) => {
            u32::from_str_radix(s.trim_start_matches("0x"), 16).map_err(de::Error::custom)
        }
    }
}

// ── de_enumerants ─────────────────────────────────────────────────────────────

fn de_enumerants<'de, D: Deserializer<'de>>(
    d: D,
) -> Result<HashMap<String, EnumerantData>, D::Error> {
    #[derive(Deserialize)]
    struct RawEnumerant {
        enumerant: String,
        #[serde(flatten)]
        data: EnumerantData,
    }
    Ok(Vec::<RawEnumerant>::deserialize(d)?
        .into_iter()
        .map(|e| (e.enumerant, e.data))
        .collect())
}

// ── types ─────────────────────────────────────────────────────────────────────

pub type CategoryKey = (Category, Kind);

#[derive(Debug, Deserialize, Eq, PartialEq, Hash)]
pub enum Category {
    BitEnum,
    ValueEnum,
    Id,
    Literal,
    Composite,
}

#[derive(Debug, Deserialize, Eq, PartialEq)]
pub struct Metadata {
    #[serde(default, deserialize_with = "de_opt_version")]
    pub version: Option<Version>,
    #[serde(rename = "lastVersion", default, deserialize_with = "de_opt_version")]
    pub last_version: Option<Version>,
    #[serde(default)] pub capabilities: Vec<String>,
    #[serde(default)] pub extensions:   Vec<String>,
    #[serde(default)] pub aliases:      Vec<String>,
    #[serde(default)] pub provisional:  bool,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct EnumerantData {
    #[serde(deserialize_with = "de_value")]
    pub value:      u32,
    #[serde(default)]
    pub parameters: Vec<Operand>,
    #[serde(flatten)]
    pub meta:       Metadata,
}

#[derive(Debug, Deserialize, Default, PartialEq)]
pub struct CategoryData {
    #[serde(default, deserialize_with = "de_enumerants")]
    pub enumerants: HashMap<String, EnumerantData>,
    #[serde(default)] pub bases: Vec<String>,
    #[serde(default)] pub doc:   Option<String>,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct InstructionData {
    pub opname:   String,
    pub class:    String,
    #[serde(default)]
    pub operands: Vec<Operand>,
    #[serde(flatten)]
    pub meta:     Metadata,
}

// ── Grammar ───────────────────────────────────────────────────────────────────

#[derive(Debug, Default, PartialEq)]
pub struct Grammar {
    pub magic_number:  u32,
    pub version:       Version,
    pub instructions:  HashMap<u32, InstructionData>,
    pub categories:    HashMap<CategoryKey, CategoryData>,
}

impl<'de> Deserialize<'de> for Grammar {
    fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        d.deserialize_map(GrammarVisitor)
    }
}

struct GrammarVisitor;

impl<'de> Visitor<'de> for GrammarVisitor {
    type Value = Grammar;

    fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "SPIR-V core grammar object")
    }

    fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Grammar, A::Error> {
        let mut magic = None::<u32>;
        let mut major = None::<u8>;
        let mut minor = None::<u8>;
        let mut revision = None::<u8>;
        let mut instructions = None::<HashMap<u32, InstructionData>>;
        let mut categories = None::<HashMap<CategoryKey, CategoryData>>;

        while let Some(key) = map.next_key::<String>()? {
            match key.as_str() {
                "magic_number" => {
                    let s = map.next_value::<String>()?;
                    magic = Some(
                        u32::from_str_radix(s.trim_start_matches("0x"), 16)
                            .map_err(de::Error::custom)?,
                    );
                }
                "major_version" => major = Some(map.next_value()?),
                "minor_version" => minor = Some(map.next_value()?),
                "revision" => revision = Some(map.next_value()?),
                "instructions" => {
                    #[derive(Deserialize)]
                    struct RawInstruction {
                        opcode: u32,
                        #[serde(flatten)]
                        data: InstructionData,
                    }
                    let raw: Vec<RawInstruction> = map.next_value()?;
                    instructions = Some(raw.into_iter().map(|i| (i.opcode, i.data)).collect());
                }
                "operand_kinds" => {
                    #[derive(Deserialize)]
                    struct RawOperandKind {
                        category: Category,
                        kind: Kind,
                        #[serde(flatten)]
                        data: CategoryData,
                    }
                    let raw: Vec<RawOperandKind> = map.next_value()?;
                    categories = Some(
                        raw.into_iter()
                            .map(|ok| ((ok.category, ok.kind), ok.data))
                            .collect(),
                    );
                }
                _ => {
                    map.next_value::<de::IgnoredAny>()?;
                }
            }
        }

        Ok(Grammar {
            magic_number: magic.ok_or_else(|| de::Error::missing_field("magic_number"))?,
            version: Version(
                major.ok_or_else(|| de::Error::missing_field("major_version"))?,
                minor.ok_or_else(|| de::Error::missing_field("minor_version"))?,
                revision.ok_or_else(|| de::Error::missing_field("revision"))?,
            ),
            instructions: instructions.ok_or_else(|| de::Error::missing_field("instructions"))?,
            categories: categories.ok_or_else(|| de::Error::missing_field("operand_kinds"))?,
        })
    }
}

// ── Kind ──────────────────────────────────────────────────────────────────────

#[derive(Debug, Deserialize, Eq, PartialEq, Hash)]
pub enum Kind {
    // BitEnum
    ImageOperands,
    FPFastMathMode,
    SelectionControl,
    LoopControl,
    FunctionControl,
    MemorySemantics,
    MemoryAccess,
    KernelProfilingInfo,
    RayFlags,
    FragmentShadingRate,
    RawAccessChainOperands,
    CooperativeMatrixOperands,
    CooperativeMatrixReduce,
    TensorAddressingOperands,
    MatrixMultiplyAccumulateOperands,
    TensorOperands,
    // ValueEnum
    SourceLanguage,
    ExecutionModel,
    AddressingModel,
    MemoryModel,
    ExecutionMode,
    StorageClass,
    Dim,
    SamplerAddressingMode,
    SamplerFilterMode,
    ImageFormat,
    ImageChannelOrder,
    ImageChannelDataType,
    FPRoundingMode,
    FPDenormMode,
    QuantizationModes,
    FPOperationMode,
    OverflowModes,
    LinkageType,
    AccessQualifier,
    HostAccessQualifier,
    FunctionParameterAttribute,
    Decoration,
    BuiltIn,
    Scope,
    GroupOperation,
    KernelEnqueueFlags,
    Capability,
    RayQueryIntersection,
    RayQueryCommittedIntersectionType,
    RayQueryCandidateIntersectionType,
    PackedVectorFormat,
    CooperativeMatrixLayout,
    CooperativeMatrixUse,
    TensorClampMode,
    InitializationModeQualifier,
    LoadCacheControl,
    StoreCacheControl,
    NamedMaximumNumberOfRegisters,
    FPEncoding,
    CooperativeVectorMatrixLayout,
    ComponentType,
    // Id
    IdResultType,
    IdResult,
    IdMemorySemantics,
    IdScope,
    IdRef,
    // Literal
    LiteralInteger,
    LiteralString,
    LiteralFloat,
    LiteralContextDependentNumber,
    LiteralExtInstInteger,
    LiteralSpecConstantOpInteger,
    // Composite
    PairLiteralIntegerIdRef,
    PairIdRefLiteralInteger,
    PairIdRefIdRef,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct Operand {
    pub kind:                   Kind,
    #[serde(default)] pub name: String,
    #[serde(default)] pub quantifier: Quantifier,
}

#[derive(Debug, Deserialize, Default, PartialEq)]
pub enum Quantifier {
    #[default]
    #[serde(rename = "")]
    Required,
    #[serde(rename = "?")]
    Optional,
    #[serde(rename = "*")]
    ZeroOrMore,
}

pub fn latest_grammar() -> Grammar {
    serde_json::from_str::<Grammar>(SPIRV_CORE_GRAMMAR).unwrap()
}
