Pumpkin/pumpkin-codegen/src/packet.rs

use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use serde::Deserialize;
use std::{collections::BTreeMap, fs};

use crate::version::JavaMinecraftVersion;

/// The newest protocol version used as the fallback for unknown versions in `PacketId::to_id`.
const LATEST_VERSION: JavaMinecraftVersion = JavaMinecraftVersion::V_26_1;

/// Raw deserialization shape for a single versioned packet mapping file.
#[derive(Deserialize)]
pub struct Packets {
    /// Protocol version number (unused at runtime; present for JSON completeness).
    #[allow(dead_code)]
    version: u32,
    /// Serverbound packet IDs grouped by protocol phase, then by packet name.
    serverbound: BTreeMap<String, BTreeMap<String, i32>>,
    /// Clientbound packet IDs grouped by protocol phase, then by packet name.
    clientbound: BTreeMap<String, BTreeMap<String, i32>>,
}

/// Generates the `TokenStream` for the `PacketId` struct, `CURRENT_MC_VERSION`, and
/// all `serverbound`/`clientbound` packet ID constants.
pub(crate) fn build() -> TokenStream {
    let assets = [
        (JavaMinecraftVersion::V_1_20_5, "1_20_5_packets.json"),
        (JavaMinecraftVersion::V_1_21, "1_21_packets.json"),
        (JavaMinecraftVersion::V_1_21_2, "1_21_2_packets.json"),
        (JavaMinecraftVersion::V_1_21_4, "1_21_4_packets.json"),
        (JavaMinecraftVersion::V_1_21_5, "1_21_5_packets.json"),
        (JavaMinecraftVersion::V_1_21_6, "1_21_6_packets.json"),
        (JavaMinecraftVersion::V_1_21_7, "1_21_7_packets.json"),
        (JavaMinecraftVersion::V_1_21_9, "1_21_9_packets.json"),
        (JavaMinecraftVersion::V_1_21_11, "1_21_11_packets.json"),
        (JavaMinecraftVersion::V_26_1, "26_1_packets.json"),
    ];

    // Parse available packet files into a BTreeMap keyed by JavaMinecraftVersion
    let mut versions = BTreeMap::new();
    for (ver, file) in assets {
        let path = format!("../assets/packet/{file}");

        let content = fs::read_to_string(&path)
            .unwrap_or_else(|_| panic!("Failed to read packet JSON file: {path}"));
        let parsed: Packets = serde_json::from_str(&content)
            .unwrap_or_else(|e| panic!("Failed to parse {path}: {e}"));

        versions.insert(ver, parsed);
    }

    // Generate PacketId struct definition and impl blocks dynamically based on versions
    let packet_id_struct = generate_struct(&versions);
    let serverbound_consts = generate_mapped_consts(&versions, true);
    let clientbound_consts = generate_mapped_consts(&versions, false);

    quote!(
        use pumpkin_util::version::JavaMinecraftVersion;

        pub const CURRENT_MC_VERSION: JavaMinecraftVersion = #LATEST_VERSION;
        pub const LOWEST_SUPPORTED_MC_VERSION: JavaMinecraftVersion = JavaMinecraftVersion::V_1_20_5;

        #packet_id_struct

        // We place the constants directly into these modules
        pub mod serverbound {
            #serverbound_consts
        }

        pub mod clientbound {
            #clientbound_consts
        }
    )
}

/// Generate the `PacketId` struct and impls (including `to_id`) dynamically based on available versions.
fn generate_struct<T>(versions: &BTreeMap<JavaMinecraftVersion, T>) -> TokenStream {
    // Build struct fields
    let mut struct_fields = TokenStream::new();
    for ver in versions.keys() {
        let ident = ver.to_field_ident();
        struct_fields.extend(quote! {
            pub #ident: i32,
        });
    }

    let latest_field_ident = LATEST_VERSION.to_field_ident();

    // Build match arms
    let mut match_arms = TokenStream::new();
    for ver in versions.keys() {
        let ident = ver.to_field_ident();
        match_arms.extend(quote! {
            #ver => self.#ident,
        });
    }

    quote! {
        #[derive(Clone, Copy, Debug)]
        pub struct PacketId {
            #struct_fields
        }

        impl PacketId {
            /// Converts the requested protocol version into the corresponding packet ID.
            /// Returns -1 if the packet does not exist in that version.
            #[must_use]
            pub const fn to_id(&self, version: JavaMinecraftVersion) -> i32 {
                #[allow(clippy::match_same_arms)]
                match version {
                    #match_arms
                    _ => self.#latest_field_ident,
                }
            }
        }

        impl PartialEq<i32> for PacketId {
            fn eq(&self, other: &i32) -> bool {
                self.#latest_field_ident == *other
            }
        }

        impl PartialEq<PacketId> for i32 {
            fn eq(&self, other: &PacketId) -> bool {
                *self == other.#latest_field_ident
            }
        }
    }
}

/// Generates `PacketId` constants for all packets in either the serverbound or clientbound direction.
///
/// # Arguments
/// - `versions` – Map from `JavaMinecraftVersion` to parsed `Packets` data.
/// - `is_serverbound` – When `true`, emits serverbound constants; otherwise emits clientbound constants.
fn generate_mapped_consts(
    versions: &BTreeMap<JavaMinecraftVersion, Packets>,
    is_serverbound: bool,
) -> TokenStream {
    let mut conv_packets = BTreeMap::<_, BTreeMap<_, _>>::new();

    for (ver, packets) in versions {
        let phases = if is_serverbound {
            &packets.serverbound
        } else {
            &packets.clientbound
        };
        for (phase, packets) in phases {
            for (name, &id) in packets {
                let sanitized_name = name.replace(['/', '-'], "_").to_uppercase();
                let const_name = format!("{}_{}", phase.to_uppercase(), sanitized_name);
                conv_packets.entry(const_name).or_default().insert(ver, id);
            }
        }
    }

    let mut output = TokenStream::new();
    for (name, values) in conv_packets {
        let mut init_pairs = TokenStream::new();
        for ver in versions.keys() {
            let id = values.get(ver).copied().unwrap_or(-1);
            let field_ident = ver.to_field_ident();
            init_pairs.extend(quote! {
                #field_ident: #id,
            });
        }
        let const_name = format_ident!("{}", name);
        output.extend(quote! {
            pub const #const_name: super::PacketId = super::PacketId {
                #init_pairs
            };
        });
    }

    output
}