use serde::Deserialize;
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Write;

#[derive(Deserialize)]
struct Config {
    #[allow(dead_code)]
    protocol_version: String,
    type_maps: BTreeMap<String, TypeMapConfig>,
}

#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct TypeMapConfig {
    #[serde(default)]
    communication_types: BTreeMap<String, u16>,
    #[serde(default)]
    data_types: BTreeMap<String, u16>,
}

struct ReservedEntry {
    name: &'static str,
    id: u16,
}

const RESERVED_COMM_TYPES: &[ReservedEntry] = &[
    ReservedEntry {
        name: "Error",
        id: 0,
    },
    ReservedEntry {
        name: "ErrorParsing",
        id: 1,
    },
    ReservedEntry {
        name: "ErrorBadVersion",
        id: 2,
    },
    ReservedEntry {
        name: "Disconnect",
        id: 3,
    },
    ReservedEntry {
        name: "Redirect",
        id: 4,
    },
    ReservedEntry {
        name: "Shutdown",
        id: 5,
    },
    ReservedEntry {
        name: "BadRequest",
        id: 6,
    },
    ReservedEntry {
        name: "Unauthorized",
        id: 7,
    },
    ReservedEntry {
        name: "Forbidden",
        id: 8,
    },
    ReservedEntry {
        name: "NotFound",
        id: 9,
    },
    ReservedEntry {
        name: "TooManyRequests",
        id: 10,
    },
    ReservedEntry {
        name: "InternalServerError",
        id: 11,
    },
    ReservedEntry {
        name: "BadGateway",
        id: 12,
    },
    ReservedEntry {
        name: "ServiceUnavailable",
        id: 13,
    },
    ReservedEntry {
        name: "GatewayTimeout",
        id: 14,
    },
    ReservedEntry {
        name: "Identification",
        id: 15,
    },
    ReservedEntry {
        name: "IdentificationResponse",
        id: 16,
    },
    ReservedEntry {
        name: "Register",
        id: 17,
    },
    ReservedEntry {
        name: "RegisterResponse",
        id: 18,
    },
    ReservedEntry {
        name: "Ping",
        id: 19,
    },
    ReservedEntry {
        name: "Pong",
        id: 20,
    },
];

const RESERVED_DATA_TYPES: &[ReservedEntry] = &[
    ReservedEntry {
        name: "Error",
        id: 0,
    },
    ReservedEntry {
        name: "ErrorParsing",
        id: 1,
    },
    ReservedEntry {
        name: "ErrorMessage",
        id: 2,
    },
    ReservedEntry {
        name: "Version",
        id: 3,
    },
    ReservedEntry {
        name: "Description",
        id: 4,
    },
    ReservedEntry {
        name: "Timestamp",
        id: 5,
    },
    ReservedEntry { name: "Id", id: 6 },
    ReservedEntry {
        name: "ClientNonce",
        id: 7,
    },
    ReservedEntry {
        name: "ServerNonce",
        id: 8,
    },
    ReservedEntry {
        name: "PublicKeys",
        id: 9,
    },
    ReservedEntry {
        name: "Signature",
        id: 10,
    },
    ReservedEntry {
        name: "Connected",
        id: 11,
    },
];

fn main() {
    let out = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap());
    let multi_version = std::env::var("CARGO_FEATURE_REGISTRY").is_ok();

    let config = match std::env::var("MTP_TYPE_MAPS") {
        Ok(config_path) => {
            let content =
                std::fs::read_to_string(&config_path).expect("Failed to read type-maps.yaml");
            serde_yaml::from_str(&content).expect("Failed to parse type-maps.yaml")
        }
        Err(_) => {
            eprintln!(
                "warning: MTP_TYPE_MAPS not set; generating types with reserved entries only"
            );
            Config {
                protocol_version: String::new(),
                type_maps: BTreeMap::new(),
            }
        }
    };

    let code = generate(&config, multi_version);
    std::fs::write(out.join("types.rs"), code).unwrap();
}

fn sorted_versions(config: &Config) -> Vec<(String, u16, u16)> {
    let mut versions: Vec<(String, u16, u16)> = config
        .type_maps
        .keys()
        .filter_map(|k| {
            let parts: Vec<&str> = k.split('.').collect();
            if parts.len() == 2 {
                let major = parts[0].parse::<u16>().ok()?;
                let minor = parts[1].parse::<u16>().ok()?;
                Some((k.clone(), major, minor))
            } else {
                None
            }
        })
        .collect();
    versions.sort_by_key(|(_, major, minor)| (*major, *minor));
    versions
}

fn generate(config: &Config, multi_version: bool) -> String {
    let mut out = String::new();
    writeln!(
        out,
        "// Auto-generated by type-map/build.rs, this file should not be edited manually"
    )
    .unwrap();
    writeln!(out).unwrap();

    let (all_comm_names, all_data_names, sorted) = if multi_version {
        let mut all_comm_names: BTreeSet<&str> = BTreeSet::new();
        let mut all_data_names: BTreeSet<&str> = BTreeSet::new();
        for (_version_key, tm) in &config.type_maps {
            for name in tm.communication_types.keys() {
                all_comm_names.insert(name.as_str());
            }
            for name in tm.data_types.keys() {
                all_data_names.insert(name.as_str());
            }
        }
        let sorted = sorted_versions(config);
        (all_comm_names, all_data_names, sorted)
    } else {
        let mut all_comm_names: BTreeSet<&str> = BTreeSet::new();
        let mut all_data_names: BTreeSet<&str> = BTreeSet::new();
        if let Some(tm) = config.type_maps.get(&config.protocol_version) {
            for name in tm.communication_types.keys() {
                all_comm_names.insert(name.as_str());
            }
            for name in tm.data_types.keys() {
                all_data_names.insert(name.as_str());
            }
        }
        (all_comm_names, all_data_names, vec![])
    };

    generate_protocol_version(&mut out, config);
    generate_comm_type_enum(&mut out, &all_comm_names);
    generate_data_type_enum(&mut out, &all_data_names);
    generate_latest_method(&mut out, config);

    if multi_version {
        generate_versioned_constructors(&mut out, &sorted);
        generate_builtin_type_maps(&mut out, &sorted);
        generate_lookup_methods(&mut out, config, &sorted);
    } else {
        generate_single_version_lookup(&mut out, config);
    }

    generate_enum_conversion_methods(&mut out);
    generate_reverse_lookups(&mut out, config, &sorted, multi_version);

    out
}

fn generate_protocol_version(out: &mut String, config: &Config) {
    let (major, minor) = if config.protocol_version.is_empty() {
        (0, 0)
    } else {
        let parts: Vec<&str> = config.protocol_version.split('.').collect();
        if parts.len() == 2 {
            let maj = parts[0].parse::<u16>().unwrap_or(0);
            let min = parts[1].parse::<u16>().unwrap_or(0);
            (maj, min)
        } else {
            (0, 0)
        }
    };
    writeln!(
        out,
        "pub const PROTOCOL_VERSION: Version = Version({}, {});",
        major, minor
    )
    .unwrap();
    writeln!(out).unwrap();
}

fn generate_latest_method(out: &mut String, config: &Config) {
    let (major, minor) = if config.protocol_version.is_empty() {
        (0, 0)
    } else {
        let parts: Vec<&str> = config.protocol_version.split('.').collect();
        if parts.len() == 2 {
            let maj = parts[0].parse::<u16>().unwrap_or(0);
            let min = parts[1].parse::<u16>().unwrap_or(0);
            (maj, min)
        } else {
            (0, 0)
        }
    };
    writeln!(out, "impl TypeMap {{").unwrap();
    writeln!(out, "    pub fn latest() -> Self {{").unwrap();
    writeln!(
        out,
        "        Self {{ version: Version({}, {}) }}",
        major, minor
    )
    .unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn generate_comm_type_enum(out: &mut String, user_names: &BTreeSet<&str>) {
    writeln!(
        out,
        "#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]"
    )
    .unwrap();
    writeln!(out, "pub enum CommunicationType {{").unwrap();

    for entry in RESERVED_COMM_TYPES {
        writeln!(out, "    {},", entry.name).unwrap();
    }
    for name in user_names {
        writeln!(out, "    {},", name).unwrap();
    }

    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "impl CommunicationType {{").unwrap();
    writeln!(out, "    pub fn name(self) -> &'static str {{").unwrap();
    writeln!(out, "        match self {{").unwrap();

    for entry in RESERVED_COMM_TYPES {
        writeln!(
            out,
            "            CommunicationType::{} => \"{}\",",
            entry.name, entry.name
        )
        .unwrap();
    }
    for name in user_names {
        writeln!(
            out,
            "            CommunicationType::{} => \"{}\",",
            name, name
        )
        .unwrap();
    }

    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "impl CommunicationType {{").unwrap();
    writeln!(out, "    pub fn from_name(s: &str) -> Option<Self> {{").unwrap();
    writeln!(out, "        match s {{").unwrap();

    for entry in RESERVED_COMM_TYPES {
        writeln!(
            out,
            "            \"{}\" => Some(CommunicationType::{}),",
            entry.name, entry.name
        )
        .unwrap();
    }
    for name in user_names {
        writeln!(
            out,
            "            \"{}\" => Some(CommunicationType::{}),",
            name, name
        )
        .unwrap();
    }

    writeln!(out, "            _ => None,").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "impl std::fmt::Display for CommunicationType {{").unwrap();
    writeln!(
        out,
        "    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {{"
    )
    .unwrap();
    writeln!(out, "        f.write_str(self.name())").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn generate_data_type_enum(out: &mut String, user_names: &BTreeSet<&str>) {
    writeln!(out, "/// Data type identifier; compile-time checked.").unwrap();
    writeln!(
        out,
        "#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]"
    )
    .unwrap();
    writeln!(out, "pub enum DataType {{").unwrap();

    for entry in RESERVED_DATA_TYPES {
        writeln!(out, "    {},", entry.name).unwrap();
    }
    for name in user_names {
        writeln!(out, "    {},", name).unwrap();
    }

    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "impl DataType {{").unwrap();
    writeln!(out, "    pub fn name(self) -> &'static str {{").unwrap();
    writeln!(out, "        match self {{").unwrap();

    for entry in RESERVED_DATA_TYPES {
        writeln!(
            out,
            "            DataType::{} => \"{}\",",
            entry.name, entry.name
        )
        .unwrap();
    }
    for name in user_names {
        writeln!(out, "            DataType::{} => \"{}\",", name, name).unwrap();
    }

    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "impl DataType {{").unwrap();
    writeln!(out, "    pub fn from_name(s: &str) -> Option<Self> {{").unwrap();
    writeln!(out, "        match s {{").unwrap();

    for entry in RESERVED_DATA_TYPES {
        writeln!(
            out,
            "            \"{}\" => Some(DataType::{}),",
            entry.name, entry.name
        )
        .unwrap();
    }
    for name in user_names {
        writeln!(
            out,
            "            \"{}\" => Some(DataType::{}),",
            name, name
        )
        .unwrap();
    }

    writeln!(out, "            _ => None,").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "impl std::fmt::Display for DataType {{").unwrap();
    writeln!(
        out,
        "    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {{"
    )
    .unwrap();
    writeln!(out, "        f.write_str(self.name())").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn generate_versioned_constructors(out: &mut String, sorted_versions: &[(String, u16, u16)]) {
    for (version_key, _major, _minor) in sorted_versions {
        let version_ident = version_key.replace('.', "_");

        writeln!(out, "impl TypeMap {{").unwrap();
        writeln!(out, "    pub fn v{}() -> Self {{", version_ident).unwrap();

        let parts: Vec<&str> = version_key.split('.').collect();
        let major = parts[0];
        let minor = parts[1];
        writeln!(
            out,
            "        Self {{ version: Version({}, {}) }}",
            major, minor
        )
        .unwrap();
        writeln!(out, "    }}").unwrap();
        writeln!(out, "}}").unwrap();
        writeln!(out).unwrap();
    }
}

fn generate_lookup_methods(
    out: &mut String,
    config: &Config,
    sorted_versions: &[(String, u16, u16)],
) {
    writeln!(out, "#[allow(unreachable_patterns, unused_variables)]").unwrap();
    writeln!(out, "impl TypeMap {{").unwrap();
    writeln!(
        out,
        "    pub fn comm_id_enum(&self, ct: CommunicationType) -> Option<u16> {{"
    )
    .unwrap();
    writeln!(out, "        match self.version {{").unwrap();
    for (version_key, major, minor) in sorted_versions {
        let tm_cfg = &config.type_maps[version_key];
        writeln!(
            out,
            "            Version({}, {}) => match ct {{",
            major, minor
        )
        .unwrap();
        for entry in RESERVED_COMM_TYPES {
            writeln!(
                out,
                "                CommunicationType::{} => Some({}),",
                entry.name, entry.id
            )
            .unwrap();
        }
        for (name, id) in &tm_cfg.communication_types {
            writeln!(
                out,
                "                CommunicationType::{} => Some({}),",
                name, id
            )
            .unwrap();
        }
        writeln!(out, "                _ => None,").unwrap();
        writeln!(out, "            }},").unwrap();
    }
    writeln!(out, "            _ => None,").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "    #[allow(unused_variables)]").unwrap();
    writeln!(
        out,
        "    pub fn data_id_enum(&self, dt: DataType) -> Option<u16> {{"
    )
    .unwrap();
    writeln!(out, "        match self.version {{").unwrap();
    for (version_key, major, minor) in sorted_versions {
        let tm_cfg = &config.type_maps[version_key];
        writeln!(
            out,
            "            Version({}, {}) => match dt {{",
            major, minor
        )
        .unwrap();
        for entry in RESERVED_DATA_TYPES {
            writeln!(
                out,
                "                DataType::{} => Some({}),",
                entry.name, entry.id
            )
            .unwrap();
        }
        for (name, id) in &tm_cfg.data_types {
            writeln!(out, "                DataType::{} => Some({}),", name, id).unwrap();
        }
        writeln!(out, "                _ => None,").unwrap();
        writeln!(out, "            }},").unwrap();
    }
    writeln!(out, "            _ => None,").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn generate_single_version_lookup(out: &mut String, config: &Config) {
    let tm_cfg = config.type_maps.get(&config.protocol_version);

    writeln!(out, "#[allow(unreachable_patterns, unused_variables)]").unwrap();
    writeln!(out, "impl TypeMap {{").unwrap();
    writeln!(
        out,
        "    pub fn comm_id_enum(&self, ct: CommunicationType) -> Option<u16> {{"
    )
    .unwrap();
    writeln!(out, "        match self.version {{").unwrap();
    writeln!(out, "            PROTOCOL_VERSION => match ct {{").unwrap();
    for entry in RESERVED_COMM_TYPES {
        writeln!(
            out,
            "                CommunicationType::{} => Some({}),",
            entry.name, entry.id
        )
        .unwrap();
    }
    if let Some(tm_cfg) = tm_cfg {
        for (name, id) in &tm_cfg.communication_types {
            writeln!(
                out,
                "                CommunicationType::{} => Some({}),",
                name, id
            )
            .unwrap();
        }
    }
    writeln!(out, "                _ => None,").unwrap();
    writeln!(out, "            }},").unwrap();
    writeln!(out, "            _ => None,").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "    #[allow(unused_variables)]").unwrap();
    writeln!(
        out,
        "    pub fn data_id_enum(&self, dt: DataType) -> Option<u16> {{"
    )
    .unwrap();
    writeln!(out, "        match self.version {{").unwrap();
    writeln!(out, "            PROTOCOL_VERSION => match dt {{").unwrap();
    for entry in RESERVED_DATA_TYPES {
        writeln!(
            out,
            "                DataType::{} => Some({}),",
            entry.name, entry.id
        )
        .unwrap();
    }
    if let Some(tm_cfg) = tm_cfg {
        for (name, id) in &tm_cfg.data_types {
            writeln!(out, "                DataType::{} => Some({}),", name, id).unwrap();
        }
    }
    writeln!(out, "                _ => None,").unwrap();
    writeln!(out, "            }},").unwrap();
    writeln!(out, "            _ => None,").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn generate_builtin_type_maps(out: &mut String, sorted_versions: &[(String, u16, u16)]) {
    writeln!(out, "pub fn builtin_type_maps() -> Vec<TypeMap> {{").unwrap();
    writeln!(out, "    vec![").unwrap();
    for (version_key, _major, _minor) in sorted_versions {
        let ident = version_key.replace('.', "_");
        writeln!(out, "        TypeMap::v{}(),", ident).unwrap();
    }
    writeln!(out, "    ]").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn generate_reverse_lookups(
    out: &mut String,
    config: &Config,
    sorted_versions: &[(String, u16, u16)],
    multi_version: bool,
) {
    let mut id_to_comm: BTreeMap<u16, String> = BTreeMap::new();
    for entry in RESERVED_COMM_TYPES {
        id_to_comm.insert(entry.id, entry.name.to_string());
    }
    if multi_version {
        for (version_key, _major, _minor) in sorted_versions {
            if let Some(tm_cfg) = config.type_maps.get(version_key) {
                for (name, id) in &tm_cfg.communication_types {
                    id_to_comm.insert(*id, name.clone());
                }
            }
        }
    } else if let Some(tm_cfg) = config.type_maps.get(&config.protocol_version) {
        for (name, id) in &tm_cfg.communication_types {
            id_to_comm.insert(*id, name.clone());
        }
    }

    writeln!(
        out,
        "pub fn communication_type_name(id: u16) -> Option<&'static str> {{"
    )
    .unwrap();
    writeln!(out, "    match id {{").unwrap();
    for (id, name) in &id_to_comm {
        writeln!(out, "        {} => Some(\"{}\"),", id, name).unwrap();
    }
    writeln!(out, "        _ => None,").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    let mut id_to_data: BTreeMap<u16, String> = BTreeMap::new();
    for entry in RESERVED_DATA_TYPES {
        id_to_data.insert(entry.id, entry.name.to_string());
    }
    if multi_version {
        for (version_key, _major, _minor) in sorted_versions {
            if let Some(tm_cfg) = config.type_maps.get(version_key) {
                for (name, id) in &tm_cfg.data_types {
                    id_to_data.insert(*id, name.clone());
                }
            }
        }
    } else if let Some(tm_cfg) = config.type_maps.get(&config.protocol_version) {
        for (name, id) in &tm_cfg.data_types {
            id_to_data.insert(*id, name.clone());
        }
    }

    writeln!(
        out,
        "pub fn data_type_name(id: u16) -> Option<&'static str> {{"
    )
    .unwrap();
    writeln!(out, "    match id {{").unwrap();
    for (id, name) in &id_to_data {
        writeln!(out, "        {} => Some(\"{}\"),", id, name).unwrap();
    }
    writeln!(out, "        _ => None,").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn generate_enum_conversion_methods(out: &mut String) {
    writeln!(out, "impl CommunicationType {{").unwrap();
    writeln!(
        out,
        "    pub fn to_id(self, tm: &TypeMap) -> CommunicationTypeId {{"
    )
    .unwrap();
    writeln!(
        out,
        "        CommunicationTypeId(tm.comm_id_enum(self).expect(\"CommunicationType not found in this TypeMap version\"))"
    )
    .unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "impl DataType {{").unwrap();
    writeln!(out, "    pub fn to_id(self, tm: &TypeMap) -> DataTypeId {{").unwrap();
    writeln!(
        out,
        "        DataTypeId(tm.data_id_enum(self).expect(\"DataType not found in this TypeMap version\"))"
    )
    .unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}