use heck::{ToShoutySnakeCase, ToUpperCamelCase};
use proc_macro2::{Span, TokenStream};
use quote::{ToTokens, format_ident, quote};
use serde::Deserialize;
use std::{
    collections::{BTreeMap, HashSet},
    fs,
};
use syn::{Ident, LitInt, LitStr};

use crate::block::BlockStateId;

/// Converts a fluid name (e.g. `water`) into its SCREAMING_SNAKE_CASE constant identifier.
fn const_fluid_name_from_fluid_name(fluid: &str) -> String {
    fluid.to_shouty_snake_case()
}

/// Derives the PascalCase property-group struct name from a fluid derived name.
fn property_group_name_from_derived_name(name: &str) -> String {
    format!("{name}_fluid_properties").to_upper_camel_case()
}

/// Maps a fluid property's original snake_case name to the generated enum type name.
struct PropertyVariantMapping {
    /// Original property name as it appears in the JSON (e.g. `"level"`).
    original_name: String,
    /// PascalCase name of the generated enum for this property (e.g. `"Level"`).
    property_enum: String,
}

/// Aggregated data for a group of fluids that share the same set of block properties.
struct PropertyCollectionData {
    /// The list of property-to-enum mappings shared by all fluids in this group.
    variant_mappings: Vec<PropertyVariantMapping>,
    /// Names of the fluids that belong to this property group.
    fluid_names: Vec<String>,
}

impl PropertyCollectionData {
    /// Appends a fluid name to this property group.
    pub fn add_fluid_name(&mut self, fluid_name: String) {
        self.fluid_names.push(fluid_name);
    }

    /// Creates a new `PropertyCollectionData` from a set of variant mappings with no fluids yet.
    pub const fn from_mappings(variant_mappings: Vec<PropertyVariantMapping>) -> Self {
        Self {
            variant_mappings,
            fluid_names: Vec::new(),
        }
    }

    /// Derives a representative name for this property group from the first fluid's name.
    pub fn derive_name(&self) -> String {
        format!("{}_like", self.fluid_names[0])
    }
}

/// A fluid block property descriptor holding the property name and its allowed values.
#[derive(Deserialize, Clone, Debug)]
pub struct PropertyStruct {
    /// Name of the generated enum (PascalCase, e.g. `"Level"`).
    pub name: String,
    /// Allowed string values for this property (e.g. `["0", "1", …, "8"]`).
    pub values: Vec<String>,
}

impl ToTokens for PropertyStruct {
    /// Emits an enum definition with `EnumVariants` impl for this property.
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let name = Ident::new(&self.name, Span::call_site());

        let variant_count = self.values.len() as u16;
        let values_index = (0..self.values.len() as u16).collect::<Vec<_>>();

        let ident_values = self.values.iter().map(|value| {
            let value_str = if value.chars().all(char::is_numeric) {
                format!("L{value}")
            } else {
                value.clone()
            };
            Ident::new(&value_str.to_upper_camel_case(), Span::call_site())
        });

        let values_2 = ident_values.clone();
        let values_3 = ident_values.clone();

        let from_values = self.values.iter().map(|value| {
            let value_str = if value.chars().all(char::is_numeric) {
                format!("L{value}")
            } else {
                value.clone()
            };
            let ident = Ident::new(&value_str.to_upper_camel_case(), Span::call_site());
            quote! {
                #value => Self::#ident
            }
        });
        let to_values = self.values.iter().map(|value| {
            let value_str = if value.chars().all(char::is_numeric) {
                format!("L{value}")
            } else {
                value.clone()
            };
            let ident = Ident::new(&value_str.to_upper_camel_case(), Span::call_site());
            quote! {
                Self::#ident => #value
            }
        });

        tokens.extend(quote! {
            #[derive(Clone, Copy, Eq, PartialEq)]
            pub enum #name {
                #(#ident_values),*
            }

            impl EnumVariants for #name {
                fn variant_count() -> u16 {
                    #variant_count
                }

                fn to_index(&self) -> u16 {
                    match self {
                        #(Self::#values_2 => #values_index),*
                    }
                }

                fn from_index(index: u16) -> Self {
                    match index {
                        #(#values_index => Self::#values_3,)*
                        _ => panic!("Invalid index: {index}"),
                    }
                }

                fn to_value(&self) -> &str {
                    match self {
                        #(#to_values),*
                    }
                }

                fn from_value(value: &str) -> Self {
                    match value {
                        #(#from_values),*,
                        _ => panic!("Invalid value: {value:?}"),
                    }
                }
            }
        });
    }
}

/// A fully resolved property group struct ready to emit as a `FluidProperties` impl.
struct FluidPropertyStruct {
    /// The underlying property collection data for this fluid group.
    data: PropertyCollectionData,
}

impl ToTokens for FluidPropertyStruct {
    /// Emits a struct definition and its `FluidProperties` impl for this property group.
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let struct_name = property_group_name_from_derived_name(&self.data.derive_name());
        let name = Ident::new(&struct_name, Span::call_site());

        let values = self.data.variant_mappings.iter().map(|entry| {
            let key = Ident::new_raw(&entry.original_name, Span::call_site());
            let value = Ident::new(&entry.property_enum, Span::call_site());

            quote! {
                #key: #value
            }
        });

        let fluid_names = &self.data.fluid_names;

        let field_names: Vec<_> = self
            .data
            .variant_mappings
            .iter()
            .rev()
            .map(|entry| Ident::new_raw(&entry.original_name, Span::call_site()))
            .collect();

        let field_types: Vec<_> = self
            .data
            .variant_mappings
            .iter()
            .rev()
            .map(|entry| Ident::new(&entry.property_enum, Span::call_site()))
            .collect();

        let to_props_values = self.data.variant_mappings.iter().map(|entry| {
            let key = &entry.original_name;
            let key2 = Ident::new_raw(&entry.original_name, Span::call_site());

            quote! {
                (#key.to_string(), self.#key2.to_value().to_string()),
            }
        });

        let from_props_values = self.data.variant_mappings.iter().map(|entry| {
            let key = &entry.original_name;
            let key2 = Ident::new_raw(&entry.original_name, Span::call_site());
            let value = Ident::new(&entry.property_enum, Span::call_site());

            quote! {
                #key => fluid_props.#key2 = #value::from_value(&value)
            }
        });

        tokens.extend(quote! {
            #[derive(Clone, Copy, Eq, PartialEq)]
            pub struct #name {
                #(pub #values),*
            }

            impl FluidProperties for #name {
                #[allow(unused_assignments)]
                fn to_index(&self) -> u16 {
                    let mut index = 0;
                    let mut multiplier = 1;

                    #(
                        index += self.#field_names.to_index() * multiplier;
                        multiplier *= #field_types::variant_count();
                    )*

                    index
                }

                #[allow(unused_assignments)]
                fn from_index(mut index: u16) -> Self {
                    Self {
                        #(
                            #field_names: {
                                let value = index % #field_types::variant_count();
                                index /= #field_types::variant_count();
                                #field_types::from_index(value)
                            }
                        ),*
                    }
                }

                fn to_state_id(&self, fluid: &Fluid) -> BlockStateId {
                    if ![#(#fluid_names),*].contains(&fluid.name) {
                        panic!("{} is not a valid fluid for {}", fluid.name, #struct_name);
                    }

                    let prop_index = self.to_index();
                    if prop_index < fluid.states.len() as u16 {
                        fluid.states[prop_index as usize].block_state_id
                    } else {
                        fluid.states[fluid.default_state_index as usize].block_state_id
                    }
                }

                fn from_state_id(id: BlockStateId, fluid: &Fluid) -> Self {
                    if ![#(#fluid_names),*].contains(&fluid.name) {
                        panic!("{} is not a valid fluid for {}", &fluid.name, #struct_name);
                    }

                    for (idx, state) in fluid.states.iter().enumerate() {
                        if state.block_state_id == id {
                            return Self::from_index(idx as u16);
                        }
                    }

                    Self::from_index(fluid.default_state_index)
                }

                fn default(fluid: &Fluid) -> Self {
                    if ![#(#fluid_names),*].contains(&fluid.name) {
                        panic!("{} is not a valid fluid for {}", &fluid.name, #struct_name);
                    }

                    Self::from_index(fluid.default_state_index)
                }

                fn to_props(&self) -> Vec<(String, String)> {
                   vec![#(#to_props_values)*]
                }

                fn from_props(props: Vec<(String, String)>, fluid: &Fluid) -> Self {
                    if ![#(#fluid_names),*].contains(&fluid.name) {
                        panic!("{} is not a valid fluid for {}", &fluid.name, #struct_name);
                    }

                    let mut fluid_props = Self::default(fluid);

                    for (key, value) in props {
                        match key.as_str() {
                            #(#from_props_values),*,
                            _ => panic!("Invalid key: {key}"),
                        }
                    }

                    fluid_props
                }
            }
        });
    }
}

/// Raw deserialization shape for a single fluid state entry from `fluids.json`.
#[derive(Deserialize, Clone)]
struct FluidState {
    /// Fraction of a full block that this fluid fills (0.0–1.0).
    height: f32,
    /// Numeric fluid level (0 = source, 1–7 = flowing).
    level: i16,
    /// Whether this state represents an empty (air) fluid slot.
    is_empty: bool,
    /// Blast resistance of the fluid in this state.
    blast_resistance: f32,
    /// Block state ID used to identify this fluid state in the world.
    block_state_id: BlockStateId,
    /// Whether the fluid is still (not flowing).
    is_still: bool,
    // We'll derive is_source and falling from existing fields instead of requiring them in JSON
}

/// A lightweight reference to a fluid state, pairing its fluid-relative index with a
/// deduplicated partial-state index.
#[derive(Clone, Debug)]
struct FluidStateRef {
    /// Index of this state within its fluid's state array.
    pub id: u16,
    /// Index into the global `FLUID_STATES` deduplication table.
    pub state_idx: u16,
}
impl ToTokens for FluidStateRef {
    /// Emits a `FluidStateRef { id, state_idx }` struct literal.
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let id = LitInt::new(&self.id.to_string(), Span::call_site());
        let state_idx = LitInt::new(&self.state_idx.to_string(), Span::call_site());
        tokens.extend(quote! {
            FluidStateRef {
                id: #id,
                state_idx: #state_idx,
            }
        });
    }
}

/// A single block-state property for a fluid, as listed in `fluids.json`.
#[derive(Deserialize, Clone)]
struct Property {
    /// Name of the property (e.g. `"level"`).
    name: String,
    /// List of valid string values for this property.
    values: Vec<String>,
}

/// Raw deserialization shape for a single fluid entry from `fluids.json`.
#[derive(Deserialize, Clone)]
pub struct Fluid {
    /// Registry name of the fluid (e.g. `"water"`).
    pub name: String,
    /// Numeric ID identifying this fluid type.
    pub id: u16,
    /// Block-state properties this fluid exposes (e.g. `level`, `falling`).
    properties: Vec<Property>,
    /// Index within `states` of the default fluid state.
    default_state_index: u16,
    /// All possible states for this fluid, one per property combination.
    states: Vec<FluidState>,
    /// Ticks between each flow-spread step (lower = faster).
    #[serde(default = "default_flow_speed")]
    flow_speed: u32,
    /// Maximum number of blocks this fluid can flow horizontally from a source.
    #[serde(default = "default_flow_distance")]
    flow_distance: u32,
    /// Whether two adjacent source blocks can create a new source block.
    #[serde(default)]
    can_convert_to_source: bool,
}

/// Default flow speed (ticks between spread steps) matching vanilla water.
const fn default_flow_speed() -> u32 {
    5 // Default to water's speed
}

/// Default horizontal flow distance in blocks, matching vanilla water.
const fn default_flow_distance() -> u32 {
    4 // Default to water's distance
}

/// Generates the `TokenStream` for the `Fluid` struct, `FluidState`, `FluidProperties` trait,
/// per-fluid property enums, and all lookup functions.
pub fn build() -> TokenStream {
    let fluids: Vec<Fluid> =
        match serde_json::from_str(&fs::read_to_string("../../assets/fluids.json").unwrap()) {
            Ok(fluids) => fluids,
            Err(e) => panic!("Failed to parse fluids.json: {e}"),
        };

    let mut constants = TokenStream::new();
    let mut id_matches = Vec::new();
    let mut type_from_name = TokenStream::new();
    let mut type_from_raw_id_arms = TokenStream::new();
    let mut fluid_from_state_id = TokenStream::new();

    // Collect from_state_id arms to sort by range width (narrower first)
    struct StateIdArm {
        start: u16,
        end: u16,
        const_name: String,
    }
    let mut state_id_arms: Vec<StateIdArm> = Vec::new();

    let mut fluid_properties_from_state_and_name = TokenStream::new();
    let mut fluid_properties_from_props_and_name = TokenStream::new();

    // Used to create property `enum`s.
    let mut property_enums: BTreeMap<String, PropertyStruct> = BTreeMap::new();
    // Property implementation for a fluid.
    let mut fluid_properties: Vec<FluidPropertyStruct> = Vec::new();
    // Mapping of a collection of property names -> fluids that have these properties.
    let mut property_collection_map: BTreeMap<Vec<String>, PropertyCollectionData> =
        BTreeMap::new();
    // Validator that we have no `enum` collisions.
    let mut enum_to_values: BTreeMap<String, Vec<String>> = BTreeMap::new();

    // Collect unique fluid states to create partial states
    let mut unique_states = Vec::new();
    let mut optimized_fluids: Vec<(String, FluidStateRef)> = Vec::new();

    for fluid in fluids {
        let id_name = LitStr::new(&fluid.name, Span::call_site());
        let const_ident = format_ident!("{}", fluid.name.to_shouty_snake_case());
        let state_id_start = fluid
            .states
            .iter()
            .map(|state| state.block_state_id.0)
            .min()
            .unwrap();
        let state_id_end = fluid
            .states
            .iter()
            .map(|state| state.block_state_id.0)
            .max()
            .unwrap();

        let id_lit = LitInt::new(&fluid.id.to_string(), Span::call_site());
        let mut properties = TokenStream::new();
        if fluid.properties.is_empty() {
            properties.extend(quote!(None));
        } else {
            let internal_properties = fluid.properties.iter().map(|property| {
                let key = LitStr::new(&property.name, Span::call_site());
                let values = property
                    .values
                    .iter()
                    .map(|value| LitStr::new(value, Span::call_site()));

                quote! {
                    (#key, &[
                        #(#values),*
                    ])
                }
            });
            properties.extend(quote! {
                Some(&[
                    #(#internal_properties),*
                ])
            });
        }

        for (idx, state) in fluid.states.iter().enumerate() {
            // Check if this state is already in `unique_states` by comparing key fields
            let already_exists = unique_states.iter().any(|s: &FluidState| {
                s.height == state.height
                    && s.level == state.level
                    && s.is_empty == state.is_empty
                    && s.blast_resistance == state.blast_resistance
                    && s.is_still == state.is_still
            });
            if !already_exists {
                unique_states.push(state.clone());
            }
            // Create a reference to the state
            let state_idx = unique_states
                .iter()
                .position(|s| {
                    s.height == state.height
                        && s.level == state.level
                        && s.is_empty == state.is_empty
                        && s.blast_resistance == state.blast_resistance
                        && s.is_still == state.is_still
                })
                .unwrap() as u16;
            optimized_fluids.push((
                fluid.name.clone(),
                FluidStateRef {
                    id: idx as u16,
                    state_idx,
                },
            ));
        }
        state_id_arms.push(StateIdArm {
            start: state_id_start,
            end: state_id_end,
            const_name: const_fluid_name_from_fluid_name(&fluid.name),
        });

        type_from_name.extend(quote! {
            #id_name => Some(&Self::#const_ident),
        });

        type_from_raw_id_arms.extend(quote! {
            #id_lit => Some(&Self::#const_ident),
        });

        let fluid_states = fluid.states.iter().map(|state| {
            let height = state.height;
            let level = state.level;
            let is_empty = state.is_empty;
            let blast_resistance = state.blast_resistance;
            let block_state_id = state.block_state_id;
            let is_still = state.is_still;
            // Derive these values based on existing fields
            let is_source = is_still; // Level 0 and still means it's a source
            let falling = false; // Default to false - we'll handle falling in the fluid behavior code

            quote! {
                FluidState {
                    height: #height,
                    level: #level,
                    is_empty: #is_empty,
                    blast_resistance: #blast_resistance,
                    block_state_id: #block_state_id,
                    is_still: #is_still,
                    is_source: #is_source,
                    falling: #falling,
                }
            }
        });
        let state_id = fluid.default_state_index;
        let flow_speed = fluid.flow_speed;
        let flow_distance = fluid.flow_distance;
        let can_convert_to_source = fluid.can_convert_to_source;

        id_matches.push(quote! {
            #id_name => Some(#id_lit),
        });

        constants.extend(quote! {
            pub const #const_ident: Fluid = Fluid {
                id: #id_lit,
                name: #id_name,
                properties: #properties,
                states: &[#(#fluid_states),*],
                default_state_index: #state_id,
                flow_speed: #flow_speed,
                flow_distance: #flow_distance,
                can_convert_to_source: #can_convert_to_source,
            };
        });

        let mut property_collection = HashSet::new();
        let mut property_mapping = Vec::new();
        for property in &fluid.properties {
            property_collection.insert(property.name.clone());

            // Get mapped property `enum` name
            let renamed_property = property.name.to_upper_camel_case();

            let expected_values = enum_to_values
                .entry(renamed_property.clone())
                .or_insert_with(|| property.values.clone());

            assert_eq!(
                expected_values, &property.values,
                "Enum overlap for '{}' ({:?} vs {:?})",
                property.name, property.values, expected_values
            );

            property_mapping.push(PropertyVariantMapping {
                original_name: property.name.clone(),
                property_enum: renamed_property.clone(),
            });

            // If this property doesn't have an `enum` yet, make one.
            let _ = property_enums
                .entry(renamed_property.clone())
                .or_insert_with(|| PropertyStruct {
                    name: renamed_property,
                    values: property.values.clone(),
                });
        }

        if !property_collection.is_empty() {
            let mut property_collection = Vec::from_iter(property_collection);
            property_collection.sort();
            property_collection_map
                .entry(property_collection)
                .or_insert_with(|| PropertyCollectionData::from_mappings(property_mapping))
                .add_fluid_name(fluid.name.clone());
        }
    }

    let unique_fluid_states = unique_states.iter().map(|state| {
        let height = state.height;
        let level = state.level;
        let is_empty = state.is_empty;
        let blast_resistance = state.blast_resistance;
        let block_state_id = state.block_state_id;
        let is_still = state.is_still;
        let is_source = is_still;
        let falling = false;
        quote! {
            PartialFluidState {
                height: #height,
                level: #level,
                is_empty: #is_empty,
                blast_resistance: #blast_resistance,
                block_state_id: #block_state_id,
                is_still: #is_still,
                is_source: #is_source,
                falling: #falling,
            }
        }
    });

    // Narrower ranges first so still fluids match before their flowing variant
    state_id_arms.sort_by_key(|arm| arm.end - arm.start);
    for arm in &state_id_arms {
        let start = LitInt::new(&arm.start.to_string(), Span::call_site());
        let end = LitInt::new(&arm.end.to_string(), Span::call_site());
        let const_ident = format_ident!("{}", arm.const_name);
        fluid_from_state_id.extend(quote! {
            #start..=#end => Some(&Fluid::#const_ident),
        });
    }

    for property_group in property_collection_map.into_values() {
        for fluid_name in &property_group.fluid_names {
            let const_fluid_name = Ident::new(
                &const_fluid_name_from_fluid_name(fluid_name),
                Span::call_site(),
            );
            let property_name = Ident::new(
                &property_group_name_from_derived_name(&property_group.derive_name()),
                Span::call_site(),
            );

            fluid_properties_from_state_and_name.extend(quote! {
                #fluid_name => Box::new(#property_name::from_state_id(id, &Fluid::#const_fluid_name)),
            });

            fluid_properties_from_props_and_name.extend(quote! {
                #fluid_name => Box::new(#property_name::from_props(props, &Fluid::#const_fluid_name)),
            });
        }

        fluid_properties.push(FluidPropertyStruct {
            data: property_group,
        });
    }

    let fluid_props = fluid_properties.iter().map(ToTokens::to_token_stream);
    let properties = property_enums.values().map(ToTokens::to_token_stream);

    quote! {
        use std::hash::{Hash, Hasher};
        use crate::tag::{Taggable, RegistryKey};
        use crate::BlockStateId;
        use pumpkin_util::resource_location::{FromResourceLocation, ResourceLocation, ToResourceLocation};

        #[derive(Clone)]
        pub struct PartialFluidState {
            pub height: f32,
            pub level: i16,
            pub is_empty: bool,
            pub blast_resistance: f32,
            pub block_state_id: BlockStateId,
            pub is_still: bool,
            pub is_source: bool,
            pub falling: bool,
        }

        #[derive(Clone)]
        pub struct FluidState {
            pub height: f32,
            pub level: i16,
            pub is_empty: bool,
            pub blast_resistance: f32,
            pub block_state_id: BlockStateId,
            pub is_still: bool,
            pub is_source: bool,
            pub falling: bool,
        }

        #[derive(Clone)]
        pub struct FluidStateRef {
            pub id: u16,
            pub state_idx: u16,
        }

        #[derive(Clone)]
        pub struct Fluid {
            pub id: u16,
            pub name: &'static str,
            pub properties: Option<&'static [(&'static str, &'static [&'static str])]>,
            pub states: &'static [FluidState],
            pub default_state_index: u16,
            pub flow_speed: u32,
            pub flow_distance: u32,
            pub can_convert_to_source: bool,
        }

        impl Hash for Fluid {
            fn hash<H: Hasher>(&self, state: &mut H) {
                self.id.hash(state);
            }
        }

        impl PartialEq for Fluid {
            fn eq(&self, other: &Self) -> bool {
                self.id == other.id
            }
        }

        impl Eq for Fluid {}

        pub const FLUID_STATES: &[PartialFluidState] = &[
            #(#unique_fluid_states),*
        ];

        pub trait EnumVariants {
            fn variant_count() -> u16;
            fn to_index(&self) -> u16;
            fn from_index(index: u16) -> Self;
            fn to_value(&self) -> &str;
            fn from_value(value: &str) -> Self;
        }

        pub trait FluidProperties where Self: 'static {
            // Convert properties to an index (`0` to `N-1`).
            fn to_index(&self) -> u16;
            // Convert an index back to properties.
            fn from_index(index: u16) -> Self where Self: Sized;

            // Convert properties to a state id.
            fn to_state_id(&self, fluid: &Fluid) -> BlockStateId;
            // Convert a state id back to properties.
            fn from_state_id(id: BlockStateId, fluid: &Fluid) -> Self where Self: Sized;
            // Get the default properties.
            fn default(fluid: &Fluid) -> Self where Self: Sized;

            // Convert properties to a `Vec` of `(name, value)`
            fn to_props(&self) -> Vec<(String, String)>;

            // Convert properties to a fluid state, and add them onto the default state.
            fn from_props(props: Vec<(String, String)>, fluid: &Fluid) -> Self where Self: Sized;
        }

        pub fn get_fluid(registry_id: &str) -> Option<&'static Fluid> {
           let key = registry_id.strip_prefix("minecraft:").unwrap_or(registry_id);
           Fluid::from_registry_key(key)
        }

        impl Fluid {
            #constants

            pub fn from_registry_key(name: &str) -> Option<&'static Self> {
                match name {
                    #type_from_name
                    _ => None
                }
            }

            pub const fn from_id(id: u16) -> Option<&'static Self> {
                match id {
                    #type_from_raw_id_arms
                    _ => None
                }
            }

            #[allow(unreachable_patterns, clippy::match_overlapping_arm)]
            pub const fn from_state_id(id: BlockStateId) -> Option<&'static Self> {
                match id.as_u16() {
                    #fluid_from_state_id
                    _ => None
                }
            }


            pub fn ident_to_fluid_id(name: &str) -> Option<u8> {
                match name {
                    #(#id_matches)*
                    _ => None
                }
            }

            #[track_caller]
            #[doc = r" Get the properties of the fluid."]
            pub fn properties(&self, id: BlockStateId) -> Box<dyn FluidProperties> {
                match self.name {
                    #fluid_properties_from_state_and_name
                    _ => panic!("Invalid state_id")
                }
            }

            #[track_caller]
            #[doc = r" Get the properties of the fluid."]
            pub fn from_properties(&self, props: Vec<(String, String)>) -> Box<dyn FluidProperties> {
                match self.name {
                    #fluid_properties_from_props_and_name
                    _ => panic!("Invalid props")
                }
            }

            pub fn same_fluid_type(a: u16, b: u16) -> bool {
                a == b
                    || (a == 1 && b == 2)
                    || (a == 2 && b == 1)
                    || (a == 3 && b == 4)
                    || (a == 4 && b == 3)
            }
            pub fn matches_type(&self, other: &Fluid) -> bool {
                Self::same_fluid_type(self.id, other.id)
            }
            pub fn to_flowing(&self) -> &'static Fluid {
                match self.id {
                    2 => &Fluid::FLOWING_WATER,
                    4 => &Fluid::FLOWING_LAVA,
                    _ => Fluid::from_id(self.id).unwrap_or(&Fluid::EMPTY),
                }
            }

            // Added helper methods for fluid behavior
            pub fn is_source(&self, state_id: BlockStateId) -> bool {
                let idx = (state_id.as_u16() as usize) % self.states.len();
                self.states[idx].is_source
            }

            pub fn is_falling(&self, state_id: BlockStateId) -> bool {
                let idx = (state_id.as_u16() as usize) % self.states.len();
                self.states[idx].falling
            }

            pub fn get_level(&self, state_id: BlockStateId) -> i16 {
                let idx = (state_id.as_u16() as usize) % self.states.len();
                self.states[idx].level
            }

            pub fn get_height(&self, state_id: BlockStateId) -> f32 {
                let idx = (state_id.as_u16() as usize) % self.states.len();
                self.states[idx].height
            }
        }

        impl ToResourceLocation for &'static Fluid {
            fn to_resource_location(&self) -> ResourceLocation {
                format!("minecraft:{}", self.name)
            }
        }

        impl FromResourceLocation for &'static Fluid {
            fn from_resource_location(resource_location: &ResourceLocation) -> Option<Self> {
                Fluid::from_registry_key(resource_location.strip_prefix("minecraft:").unwrap_or(resource_location))
            }
        }

        impl FluidStateRef {
            pub fn get_state(&self) -> FluidState {
                let partial_state = &FLUID_STATES[self.state_idx as usize];
                FluidState {
                    height: partial_state.height,
                    level: partial_state.level,
                    is_empty: partial_state.is_empty,
                    blast_resistance: partial_state.blast_resistance,
                    block_state_id: partial_state.block_state_id,
                    is_still: partial_state.is_still,
                    is_source: partial_state.is_source,
                    falling: partial_state.falling,
                }
            }
        }

        impl Taggable for Fluid {
            #[inline]
            fn tag_key() -> RegistryKey {
                RegistryKey::Fluid
            }

            #[inline]
            fn registry_key(&self) -> &str {
                self.name
            }

            #[inline]
            fn registry_id(&self) -> u16 {
                self.id
            }
        }

        // Added FluidLevel enum and required constants
        pub const FLUID_LEVEL_SOURCE: i32 = 0;
        pub const FLUID_LEVEL_FLOWING_MAX: i32 = 8;
        pub const FLUID_MIN_HEIGHT: f32 = 0.0;
        pub const FLUID_MAX_HEIGHT: f32 = 1.0;

        #(#properties)*

        #(#fluid_props)*
    }
}
