#!/usr/bin/env python3
# pyright: reportMissingImports=false, reportExplicitAny=false, reportAny=false, reportUnknownVariableType=false, reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownParameterType=false, reportUnknownLambdaType=false, reportUnusedCallResult=false
"""Generate Rust and Python service enums from defs/bloomberg.toml."""

from __future__ import annotations

import argparse
from pathlib import Path
import shutil
import subprocess
import sys
from typing import Any

try:
    import tomllib
except ModuleNotFoundError:  # pragma: no cover - Python 3.10 fallback
    tomllib = __import__("tomli")


ROOT = Path(__file__).resolve().parents[1]
DEFS_PATH = ROOT / "defs" / "bloomberg.toml"
RUST_OUT = ROOT / "crates" / "xbbg-async" / "src" / "services.rs"
PY_OUT = ROOT / "py-xbbg" / "src" / "xbbg" / "_services_gen.py"

RUST_HEADER = "// @generated by codegen/generate.py from defs/bloomberg.toml — DO NOT EDIT"
PY_HEADER = "# @generated by codegen/generate.py from defs/bloomberg.toml — DO NOT EDIT"


def to_pascal(key: str) -> str:
    """Convert UPPER_SNAKE to PascalCase."""
    return "".join(part.capitalize() for part in key.split("_") if part)


def sorted_items(table: dict[str, Any]) -> list[tuple[str, Any]]:
    return sorted(table.items(), key=lambda pair: pair[0])


def _render_python_enum(
    lines: list[str],
    *,
    class_name: str,
    class_doc: str,
    table: dict[str, dict[str, Any]],
    value_key: str,
) -> None:
    lines.append(f"class {class_name}(str, Enum):")
    lines.append(f'    """{class_doc}"""')
    lines.append("")
    for key, data in sorted_items(table):
        lines.append(f'    {key} = "{escape_python_string(data[value_key])}"')
        lines.append(f'    """{escape_python_doc(data["doc"])}"""')
        lines.append("")
    if lines[-1] == "":
        _ = lines.pop()
    lines.append("")
    lines.append("")


def rust_name(key: str, data: dict[str, Any]) -> str:
    value = data.get("rust_name")
    if isinstance(value, str) and value:
        return value
    return to_pascal(key)


def escape_rust_string(value: str) -> str:
    return value.replace("\\", "\\\\").replace('"', '\\"')


def escape_python_string(value: str) -> str:
    return value.replace("\\", "\\\\").replace('"', '\\"')


def escape_python_doc(value: str) -> str:
    return value.replace('"""', '\\"\\"\\"')


def load_defs(path: Path) -> dict[str, Any]:
    with path.open("rb") as handle:
        data = tomllib.load(handle)
    required = ["services", "operations", "extractors", "formats", "output_modes"]
    for key in required:
        if key not in data:
            raise ValueError(f"Missing [{key}] section in {path}")
    return data


def find_default_name(section: str, table: dict[str, dict[str, Any]], names: dict[str, str]) -> str:
    defaults = [key for key, data in table.items() if data.get("default") is True]
    if len(defaults) != 1:
        raise ValueError(f"Section [{section}] must define exactly one default item")
    return names[defaults[0]]


def validate_refs(defs: dict[str, Any]) -> None:
    service_keys = set(defs["services"].keys())
    extractor_keys = set(defs["extractors"].keys())

    for op_key, op in defs["operations"].items():
        service_key = op.get("service")
        extractor_key = op.get("extractor")
        if service_key is not None and service_key not in service_keys:
            raise ValueError(f"Operation {op_key} references unknown service {service_key}")
        if extractor_key is not None and extractor_key not in extractor_keys:
            raise ValueError(f"Operation {op_key} references unknown extractor {extractor_key}")


def render_rust(defs: dict[str, Any]) -> str:
    services: dict[str, dict[str, Any]] = defs["services"]
    operations: dict[str, dict[str, Any]] = defs["operations"]
    extractors: dict[str, dict[str, Any]] = defs["extractors"]
    formats: dict[str, dict[str, Any]] = defs["formats"]
    output_modes: dict[str, dict[str, Any]] = defs["output_modes"]

    service_names = {key: rust_name(key, data) for key, data in services.items()}
    operation_names = {key: rust_name(key, data) for key, data in operations.items()}
    extractor_names = {key: rust_name(key, data) for key, data in extractors.items()}
    format_names = {key: rust_name(key, data) for key, data in formats.items()}
    output_mode_names = {key: rust_name(key, data) for key, data in output_modes.items()}

    default_extractor_name = find_default_name("extractors", extractors, extractor_names)
    default_format_name = find_default_name("formats", formats, format_names)
    default_output_mode_name = find_default_name("output_modes", output_modes, output_mode_names)

    lines: list[str] = []
    lines.append(RUST_HEADER)
    lines.append("")
    lines.append("use std::convert::Infallible;")
    lines.append("use std::fmt;")
    lines.append("use std::str::FromStr;")
    lines.append("")
    lines.append("use serde::{Deserialize, Serialize};")
    lines.append("")
    lines.append("/// Bloomberg service URIs.")
    lines.append("///")
    lines.append("/// Standard Bloomberg API services with URIs from the Bloomberg C++ SDK.")
    lines.append("/// Use [Service::Custom] for services not listed here.")
    lines.append("#[derive(Clone, Debug, PartialEq, Eq, Hash)]")
    lines.append("pub enum Service {")
    for key, data in sorted_items(services):
        lines.append("    /// {}".format(data["doc"]))
        lines.append(f"    {service_names[key]},")
    lines.append("    /// Custom service URI not listed above.")
    lines.append("    Custom(String),")
    lines.append("}")
    lines.append("")
    lines.append("impl Service {")
    lines.append("    /// Returns the Bloomberg service URI string.")
    lines.append("    pub fn as_str(&self) -> &str {")
    lines.append("        match self {")
    for key, data in sorted_items(services):
        lines.append(
            '            Self::{} => "{}",'.format(
                service_names[key],
                escape_rust_string(data["uri"]),
            )
        )
    lines.append("            Self::Custom(s) => s.as_str(),")
    lines.append("        }")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl fmt::Display for Service {")
    lines.append("    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {")
    lines.append('        write!(f, "{}", self.as_str())')
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl FromStr for Service {")
    lines.append("    type Err = Infallible;")
    lines.append("")
    lines.append("    fn from_str(s: &str) -> Result<Self, Self::Err> {")
    lines.append("        let svc = match s {")
    for key, data in sorted_items(services):
        lines.append(
            '            "{}" => Self::{},'.format(
                escape_rust_string(data["uri"]),
                service_names[key],
            )
        )
    lines.append("            other => Self::Custom(other.to_string()),")
    lines.append("        };")
    lines.append("        Ok(svc)")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl Serialize for Service {")
    lines.append("    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>")
    lines.append("    where")
    lines.append("        S: serde::Serializer,")
    lines.append("    {")
    lines.append("        serializer.serialize_str(self.as_str())")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl<'de> Deserialize<'de> for Service {")
    lines.append("    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>")
    lines.append("    where")
    lines.append("        D: serde::Deserializer<'de>,")
    lines.append("    {")
    lines.append("        let s = String::deserialize(deserializer)?;")
    lines.append("        Ok(match Self::from_str(&s) {")
    lines.append("            Ok(service) => service,")
    lines.append("            Err(never) => match never {},")
    lines.append("        })")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("/// Bloomberg request operation names.")
    lines.append("///")
    lines.append("/// Standard Bloomberg API operation types from the Bloomberg C++ SDK.")
    lines.append("/// Use [Operation::Custom] for operations not listed here.")
    lines.append("#[derive(Clone, Debug, PartialEq, Eq, Hash)]")
    lines.append("pub enum Operation {")
    for key, data in sorted_items(operations):
        lines.append("    /// {}".format(data["doc"]))
        lines.append(f"    {operation_names[key]},")
    lines.append("    /// Custom operation name not listed above.")
    lines.append("    Custom(String),")
    lines.append("}")
    lines.append("")
    lines.append("impl Operation {")
    lines.append("    /// Returns the Bloomberg operation name string.")
    lines.append("    pub fn as_str(&self) -> &str {")
    lines.append("        match self {")
    for key, data in sorted_items(operations):
        lines.append(
            '            Self::{} => "{}",'.format(
                operation_names[key],
                escape_rust_string(data["value"]),
            )
        )
    lines.append("            Self::Custom(s) => s.as_str(),")
    lines.append("        }")
    lines.append("    }")
    lines.append("")
    lines.append("    /// Returns the default extractor for this operation.")
    lines.append("    pub fn default_extractor(&self) -> ExtractorType {")
    lines.append("        match self {")
    for key, data in sorted_items(operations):
        extractor_key = data.get("extractor")
        if extractor_key is None:
            extractor_variant = "Generic"
        else:
            extractor_variant = extractor_names[extractor_key]
        lines.append(f"            Self::{operation_names[key]} => ExtractorType::{extractor_variant},")
    lines.append("            Self::Custom(_) => ExtractorType::Generic,")
    lines.append("        }")
    lines.append("    }")
    lines.append("")
    lines.append("    /// Returns the default service for this operation, if known.")
    lines.append("    pub fn default_service(&self) -> Option<Service> {")
    lines.append("        match self {")
    for key, data in sorted_items(operations):
        service_key = data.get("service")
        if service_key is None:
            service_expr = "None"
        else:
            service_expr = f"Some(Service::{service_names[service_key]})"
        lines.append(f"            Self::{operation_names[key]} => {service_expr},")
    lines.append("            Self::Custom(_) => None,")
    lines.append("        }")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl fmt::Display for Operation {")
    lines.append("    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {")
    lines.append('        write!(f, "{}", self.as_str())')
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl FromStr for Operation {")
    lines.append("    type Err = Infallible;")
    lines.append("")
    lines.append("    fn from_str(s: &str) -> Result<Self, Self::Err> {")
    lines.append("        let op = match s {")
    for key, data in sorted_items(operations):
        lines.append(
            '            "{}" => Self::{},'.format(
                escape_rust_string(data["value"]),
                operation_names[key],
            )
        )
    lines.append("            other => Self::Custom(other.to_string()),")
    lines.append("        };")
    lines.append("        Ok(op)")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl Serialize for Operation {")
    lines.append("    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>")
    lines.append("    where")
    lines.append("        S: serde::Serializer,")
    lines.append("    {")
    lines.append("        serializer.serialize_str(self.as_str())")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl<'de> Deserialize<'de> for Operation {")
    lines.append("    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>")
    lines.append("    where")
    lines.append("        D: serde::Deserializer<'de>,")
    lines.append("    {")
    lines.append("        let s = String::deserialize(deserializer)?;")
    lines.append("        Ok(match Self::from_str(&s) {")
    lines.append("            Ok(operation) => operation,")
    lines.append("            Err(never) => match never {},")
    lines.append("        })")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("/// Extractor type hint for Arrow conversion.")
    lines.append("///")
    lines.append("/// Tells the pump which Arrow schema/extractor to use for the response.")
    lines.append("#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]")
    lines.append("pub enum ExtractorType {")
    for key, data in sorted_items(extractors):
        lines.append("    /// {}".format(data["doc"]))
        if extractor_names[key] == default_extractor_name:
            lines.append("    #[default]")
        lines.append(f"    {extractor_names[key]},")
    lines.append("}")
    lines.append("")
    lines.append("impl ExtractorType {")
    lines.append("    /// Parse extractor type from string (from Python).")
    lines.append("    pub fn parse(s: &str) -> Option<Self> {")
    lines.append("        match s {")
    for key, data in sorted_items(extractors):
        lines.append(
            '            "{}" => Some(Self::{}),'.format(
                escape_rust_string(data["value"]),
                extractor_names[key],
            )
        )
    lines.append("            _ => None,")
    lines.append("        }")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("/// Output format for reference data (bdp/bdh).")
    lines.append("///")
    lines.append("/// Controls the shape and typing of the output Arrow table.")
    lines.append("#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]")
    lines.append('#[serde(rename_all = "snake_case")]')
    lines.append("pub enum Format {")
    for key, data in sorted_items(formats):
        lines.append("    /// {}".format(data["doc"]))
        if format_names[key] == default_format_name:
            lines.append("    #[default]")
        lines.append(f"    {format_names[key]},")
    lines.append("}")
    lines.append("")
    lines.append("impl Format {")
    lines.append("    /// Returns the format identifier string.")
    lines.append("    pub fn as_str(&self) -> &str {")
    lines.append("        match self {")
    for key, data in sorted_items(formats):
        lines.append(
            '            Self::{} => "{}",'.format(
                format_names[key],
                escape_rust_string(data["value"]),
            )
        )
    lines.append("        }")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl fmt::Display for Format {")
    lines.append("    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {")
    lines.append('        write!(f, "{}", self.as_str())')
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl FromStr for Format {")
    lines.append("    type Err = String;")
    lines.append("")
    lines.append("    fn from_str(s: &str) -> Result<Self, Self::Err> {")
    lines.append("        match s {")
    for key, data in sorted_items(formats):
        lines.append(
            '            "{}" => Ok(Self::{}),'.format(
                escape_rust_string(data["value"]),
                format_names[key],
            )
        )
    lines.append('            other => Err(format!("Unknown format: {}", other)),')
    lines.append("        }")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("/// Output mode for generic requests.")
    lines.append("///")
    lines.append("/// Controls how Bloomberg responses are converted before returning.")
    lines.append("#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]")
    lines.append('#[serde(rename_all = "snake_case")]')
    lines.append("pub enum OutputMode {")
    for key, data in sorted_items(output_modes):
        lines.append("    /// {}".format(data["doc"]))
        if output_mode_names[key] == default_output_mode_name:
            lines.append("    #[default]")
        lines.append(f"    {output_mode_names[key]},")
    lines.append("}")
    lines.append("")
    lines.append("impl OutputMode {")
    lines.append("    /// Returns the output mode identifier string.")
    lines.append("    pub fn as_str(&self) -> &str {")
    lines.append("        match self {")
    for key, data in sorted_items(output_modes):
        lines.append(
            '            Self::{} => "{}",'.format(
                output_mode_names[key],
                escape_rust_string(data["value"]),
            )
        )
    lines.append("        }")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl fmt::Display for OutputMode {")
    lines.append("    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {")
    lines.append('        write!(f, "{}", self.as_str())')
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("impl FromStr for OutputMode {")
    lines.append("    type Err = String;")
    lines.append("")
    lines.append("    fn from_str(s: &str) -> Result<Self, Self::Err> {")
    lines.append("        match s {")
    for key, data in sorted_items(output_modes):
        lines.append(
            '            "{}" => Ok(Self::{}),'.format(
                escape_rust_string(data["value"]),
                output_mode_names[key],
            )
        )
    lines.append('            other => Err(format!("Unknown output mode: {}", other)),')
    lines.append("        }")
    lines.append("    }")
    lines.append("}")
    lines.append("")
    lines.append("#[cfg(test)]")
    lines.append("mod tests {")
    lines.append("    use super::*;")
    lines.append("")

    roundtrip_service_key = "REFDATA" if "REFDATA" in services else sorted(services.keys())[0]
    roundtrip_service_uri = services[roundtrip_service_key]["uri"]
    roundtrip_service_variant = service_names[roundtrip_service_key]

    lines.append("    #[test]")
    lines.append("    fn service_roundtrip() {")
    lines.append(
        f'        assert_eq!(Service::from_str("{escape_rust_string(roundtrip_service_uri)}").unwrap(), Service::{roundtrip_service_variant});'
    )
    lines.append(
        f'        assert_eq!(Service::{roundtrip_service_variant}.as_str(), "{escape_rust_string(roundtrip_service_uri)}");'
    )
    lines.append(
        f'        assert_eq!(Service::{roundtrip_service_variant}.to_string(), "{escape_rust_string(roundtrip_service_uri)}");'
    )
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn service_custom() {")
    lines.append('        let custom = Service::from_str("//blp/custom").unwrap();')
    lines.append('        assert_eq!(custom, Service::Custom("//blp/custom".to_string()));')
    lines.append('        assert_eq!(custom.as_str(), "//blp/custom");')
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn service_all_known() {")
    lines.append("        for (uri, expected) in [")
    for key, data in sorted_items(services):
        lines.append(
            '            ("{}", Service::{}),'.format(
                escape_rust_string(data["uri"]),
                service_names[key],
            )
        )
    lines.append("        ] {")
    lines.append("            assert_eq!(Service::from_str(uri).unwrap(), expected);")
    lines.append("            assert_eq!(expected.as_str(), uri);")
    lines.append("        }")
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn service_serde() {")
    lines.append(f"        let svc = Service::{roundtrip_service_variant};")
    lines.append("        let json = serde_json::to_string(&svc).unwrap();")
    lines.append(f'        assert_eq!(json, r#""{escape_rust_string(roundtrip_service_uri)}""#);')
    lines.append("        let back: Service = serde_json::from_str(&json).unwrap();")
    lines.append(f"        assert_eq!(back, Service::{roundtrip_service_variant});")
    lines.append("    }")
    lines.append("")

    roundtrip_operation_key = "REFERENCE_DATA" if "REFERENCE_DATA" in operations else sorted(operations.keys())[0]
    roundtrip_operation_value = operations[roundtrip_operation_key]["value"]
    roundtrip_operation_variant = operation_names[roundtrip_operation_key]
    serde_operation_key = "HISTORICAL_DATA" if "HISTORICAL_DATA" in operations else roundtrip_operation_key
    serde_operation_value = operations[serde_operation_key]["value"]
    serde_operation_variant = operation_names[serde_operation_key]

    lines.append("    #[test]")
    lines.append("    fn operation_roundtrip() {")
    lines.append(
        f'        assert_eq!(Operation::from_str("{escape_rust_string(roundtrip_operation_value)}").unwrap(), Operation::{roundtrip_operation_variant});'
    )
    lines.append(
        f'        assert_eq!(Operation::{roundtrip_operation_variant}.as_str(), "{escape_rust_string(roundtrip_operation_value)}");'
    )
    lines.append(
        f'        assert_eq!(Operation::{roundtrip_operation_variant}.to_string(), "{escape_rust_string(roundtrip_operation_value)}");'
    )
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn operation_custom() {")
    lines.append('        let custom = Operation::from_str("CustomRequest").unwrap();')
    lines.append('        assert_eq!(custom, Operation::Custom("CustomRequest".to_string()));')
    lines.append('        assert_eq!(custom.as_str(), "CustomRequest");')
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn operation_all_known() {")
    lines.append("        for (op, expected) in [")
    for key, data in sorted_items(operations):
        lines.append(
            '            ("{}", Operation::{}),'.format(
                escape_rust_string(data["value"]),
                operation_names[key],
            )
        )
    lines.append("        ] {")
    lines.append("            assert_eq!(Operation::from_str(op).unwrap(), expected);")
    lines.append("            assert_eq!(expected.as_str(), op);")
    lines.append("        }")
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn operation_serde() {")
    lines.append(f"        let op = Operation::{serde_operation_variant};")
    lines.append("        let json = serde_json::to_string(&op).unwrap();")
    lines.append(f'        assert_eq!(json, r#""{escape_rust_string(serde_operation_value)}""#);')
    lines.append("        let back: Operation = serde_json::from_str(&json).unwrap();")
    lines.append(f"        assert_eq!(back, Operation::{serde_operation_variant});")
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn operation_default_extractor() {")
    for key, data in sorted_items(operations):
        extractor_key = data.get("extractor")
        if extractor_key is None:
            extractor_variant = "Generic"
        else:
            extractor_variant = extractor_names[extractor_key]
        lines.append(
            f"        assert_eq!(Operation::{operation_names[key]}.default_extractor(), ExtractorType::{extractor_variant});"
        )
    lines.append('        assert_eq!(Operation::Custom("foo".into()).default_extractor(), ExtractorType::Generic);')
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn operation_default_service() {")
    for key, data in sorted_items(operations):
        service_key = data.get("service")
        if service_key is None:
            service_expr = "None"
        else:
            service_expr = f"Some(Service::{service_names[service_key]})"
        lines.append(f"        assert_eq!(Operation::{operation_names[key]}.default_service(), {service_expr});")
    lines.append('        assert_eq!(Operation::Custom("foo".into()).default_service(), None);')
    lines.append("    }")
    lines.append("")
    lowercase_ops = []
    for key, data in sorted_items(operations):
        value = data["value"]
        if value and value[0].islower():
            lowercase_ops.append((operation_names[key], value))

    lines.append("    #[test]")
    lines.append("    fn operation_lowercase_operations() {")
    if lowercase_ops:
        for variant, value in lowercase_ops:
            lines.append(
                f'        assert_eq!(Operation::from_str("{escape_rust_string(value)}").unwrap(), Operation::{variant});'
            )
            lines.append(f'        assert_eq!(Operation::{variant}.as_str(), "{escape_rust_string(value)}");')
            lines.append("")
        if lines[-1] == "":
            _ = lines.pop()
    else:
        lines.append("        assert!(true);")
    lines.append("    }")
    lines.append("")

    roundtrip_format_key = "LONG" if "LONG" in formats else sorted(formats.keys())[0]
    roundtrip_format_variant = format_names[roundtrip_format_key]
    roundtrip_format_value = formats[roundtrip_format_key]["value"]
    serde_format_key = "LONG_TYPED" if "LONG_TYPED" in formats else roundtrip_format_key
    serde_format_variant = format_names[serde_format_key]
    serde_format_value = formats[serde_format_key]["value"]

    lines.append("    #[test]")
    lines.append("    fn format_roundtrip() {")
    lines.append(
        f'        assert_eq!(Format::from_str("{escape_rust_string(roundtrip_format_value)}").unwrap(), Format::{roundtrip_format_variant});'
    )
    lines.append(
        f'        assert_eq!(Format::{roundtrip_format_variant}.as_str(), "{escape_rust_string(roundtrip_format_value)}");'
    )
    lines.append(
        f'        assert_eq!(Format::{roundtrip_format_variant}.to_string(), "{escape_rust_string(roundtrip_format_value)}");'
    )
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn format_all_variants() {")
    lines.append("        for (s, expected) in [")
    for key, data in sorted_items(formats):
        lines.append(
            '            ("{}", Format::{}),'.format(
                escape_rust_string(data["value"]),
                format_names[key],
            )
        )
    lines.append("        ] {")
    lines.append("            assert_eq!(Format::from_str(s).unwrap(), expected);")
    lines.append("            assert_eq!(expected.as_str(), s);")
    lines.append("        }")
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn format_invalid() {")
    lines.append('        assert!(Format::from_str("invalid").is_err());')
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn format_serde() {")
    lines.append(f"        let fmt = Format::{serde_format_variant};")
    lines.append("        let json = serde_json::to_string(&fmt).unwrap();")
    lines.append(f'        assert_eq!(json, r#""{escape_rust_string(serde_format_value)}""#);')
    lines.append("        let back: Format = serde_json::from_str(&json).unwrap();")
    lines.append(f"        assert_eq!(back, Format::{serde_format_variant});")
    lines.append("    }")
    lines.append("")

    roundtrip_mode_key = "ARROW" if "ARROW" in output_modes else sorted(output_modes.keys())[0]
    roundtrip_mode_variant = output_mode_names[roundtrip_mode_key]
    roundtrip_mode_value = output_modes[roundtrip_mode_key]["value"]
    serde_mode_key = "JSON" if "JSON" in output_modes else roundtrip_mode_key
    serde_mode_variant = output_mode_names[serde_mode_key]
    serde_mode_value = output_modes[serde_mode_key]["value"]

    lines.append("    #[test]")
    lines.append("    fn output_mode_roundtrip() {")
    lines.append(
        f'        assert_eq!(OutputMode::from_str("{escape_rust_string(roundtrip_mode_value)}").unwrap(), OutputMode::{roundtrip_mode_variant});'
    )
    lines.append(
        f'        assert_eq!(OutputMode::{roundtrip_mode_variant}.as_str(), "{escape_rust_string(roundtrip_mode_value)}");'
    )
    lines.append(
        f'        assert_eq!(OutputMode::{roundtrip_mode_variant}.to_string(), "{escape_rust_string(roundtrip_mode_value)}");'
    )
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn output_mode_all_variants() {")
    lines.append("        for (s, expected) in [")
    for key, data in sorted_items(output_modes):
        lines.append(
            '            ("{}", OutputMode::{}),'.format(
                escape_rust_string(data["value"]),
                output_mode_names[key],
            )
        )
    lines.append("        ] {")
    lines.append("            assert_eq!(OutputMode::from_str(s).unwrap(), expected);")
    lines.append("            assert_eq!(expected.as_str(), s);")
    lines.append("        }")
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn output_mode_invalid() {")
    lines.append('        assert!(OutputMode::from_str("invalid").is_err());')
    lines.append("    }")
    lines.append("")
    lines.append("    #[test]")
    lines.append("    fn output_mode_serde() {")
    lines.append(f"        let mode = OutputMode::{serde_mode_variant};")
    lines.append("        let json = serde_json::to_string(&mode).unwrap();")
    lines.append(f'        assert_eq!(json, r#""{escape_rust_string(serde_mode_value)}""#);')
    lines.append("        let back: OutputMode = serde_json::from_str(&json).unwrap();")
    lines.append(f"        assert_eq!(back, OutputMode::{serde_mode_variant});")
    lines.append("    }")
    lines.append("}")

    return "\n".join(lines) + "\n"


def render_python(defs: dict[str, Any]) -> str:
    services: dict[str, dict[str, Any]] = defs["services"]
    operations: dict[str, dict[str, Any]] = defs["operations"]
    extractors: dict[str, dict[str, Any]] = defs["extractors"]
    formats: dict[str, dict[str, Any]] = defs["formats"]
    output_modes: dict[str, dict[str, Any]] = defs["output_modes"]

    lines: list[str] = []
    lines.append(PY_HEADER)
    lines.append("")
    lines.append("from __future__ import annotations")
    lines.append("")
    lines.append("from enum import Enum")
    lines.append("")
    lines.append("")
    _render_python_enum(
        lines,
        class_name="Service",
        class_doc="Bloomberg service URIs.",
        table=services,
        value_key="uri",
    )
    _render_python_enum(
        lines,
        class_name="Operation",
        class_doc="Bloomberg request operation names.",
        table=operations,
        value_key="value",
    )
    _render_python_enum(
        lines,
        class_name="ExtractorHint",
        class_doc="Hint for which Arrow extractor to use.",
        table=extractors,
        value_key="value",
    )
    _render_python_enum(
        lines,
        class_name="Format",
        class_doc="Output format for reference data.",
        table=formats,
        value_key="value",
    )
    _render_python_enum(
        lines,
        class_name="OutputMode",
        class_doc="Output mode for generic requests.",
        table=output_modes,
        value_key="value",
    )

    while lines and lines[-1] == "":
        _ = lines.pop()
    return "\n".join(lines) + "\n"


def is_stale(path: Path, expected: str) -> bool:
    if not path.exists():
        return True
    actual = path.read_text(encoding="utf-8")
    return actual != expected


def write_if_changed(path: Path, expected: str) -> bool:
    if not is_stale(path, expected):
        return False
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(expected, encoding="utf-8", newline="\n")
    return True


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate service enums from defs/bloomberg.toml")
    parser.add_argument(
        "--check",
        action="store_true",
        help="Check whether generated files are up to date (CI mode)",
    )
    return parser.parse_args(argv)


def _rustfmt(text: str) -> str:
    """Run rustfmt on generated Rust code so it matches cargo fmt output."""
    rustfmt = shutil.which("rustfmt")
    if rustfmt is None:
        return text
    result = subprocess.run(
        [rustfmt, "--edition", "2021"],
        input=text,
        capture_output=True,
        text=True,
        encoding="utf-8",
        check=False,
    )
    return result.stdout if result.returncode == 0 else text


def main(argv: list[str]) -> int:
    args = parse_args(argv)

    defs = load_defs(DEFS_PATH)
    validate_refs(defs)
    rust_text = _rustfmt(render_rust(defs))
    python_text = render_python(defs)

    targets = [
        (RUST_OUT, rust_text),
        (PY_OUT, python_text),
    ]

    if args.check:
        stale: list[Path] = []
        for path, expected in targets:
            if is_stale(path, expected):
                stale.append(path)
        if stale:
            for path in stale:
                print(f"Stale generated file: {path.relative_to(ROOT)}")
            return 1
        print("Generated files are up to date.")
        return 0

    changed_any = False
    for path, expected in targets:
        changed = write_if_changed(path, expected)
        if changed:
            changed_any = True
            print(f"Updated {path.relative_to(ROOT)}")
        else:
            print(f"Unchanged {path.relative_to(ROOT)}")

    if not changed_any:
        print("No changes needed.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
