cf-colo-hint 0.1.6

#!/usr/bin/env python3
"""
Code generator for cloudflare-locations-rs.

Reads components.json (Cloudflare status page data) and where.durableobjects.live.json
(DO region mapping data) to generate Rust code with colo constants and region mappings.

Usage:
    python3 codegen.py

This will regenerate src/generated.rs with the latest data.
"""

import json
import re
from pathlib import Path

# Location hints supported by Durable Objects
LOCATION_HINTS = {
    "wnam": ("WNam", "Western North America"),
    "enam": ("ENam", "Eastern North America"),
    "sam": ("Sam", "South America"),
    "weur": ("WEur", "Western Europe"),
    "eeur": ("EEur", "Eastern Europe"),
    "apac": ("Apac", "Asia-Pacific"),
    "oc": ("Oc", "Oceania"),
    "afr": ("Afr", "Africa"),
    "me": ("Me", "Middle East"),
}

# Cloudflare status page group IDs to region names
CF_GROUPS = {
    "00gpj4s37mz4": ("Africa", "afr"),
    "77867vxkttgw": ("Asia", "apac"),
    "zqxhg7y54vy8": ("Europe", "weur"),  # Default to weur, will be overridden by DO data
    "91blz4ztt7dm": ("LatinAmericaCaribbean", "sam"),
    "m3639x4txd08": ("MiddleEast", "me"),
    "4l01sk5cdn5c": ("NorthAmerica", "wnam"),  # Default to wnam, will be overridden
    "q6qm6fvkst4h": ("Oceania", "oc"),
}


def load_json(path: Path) -> dict:
    with open(path) as f:
        return json.load(f)


def extract_colos_from_components(components_data: dict) -> dict[str, dict]:
    """Extract colo information from components.json."""
    colos = {}
    for c in components_data["components"]:
        if c.get("group") is False and c.get("group_id"):
            match = re.search(r"\(([A-Z0-9]{3})\)$", c["name"])
            if match:
                code = match.group(1)
                name_part = c["name"].rsplit(" - ", 1)[0].strip()
                group_id = c["group_id"]
                group_info = CF_GROUPS.get(group_id, ("Unknown", None))
                colos[code] = {
                    "name": name_part,
                    "cf_region": group_info[0],
                    "cf_region_hint": group_info[1],
                }
    return colos


def merge_do_data(colos: dict[str, dict], do_data: dict) -> dict[str, dict]:
    """Merge Durable Objects region data into colo info."""
    for code, do_info in do_data.get("colos", {}).items():
        if code in colos:
            colos[code]["nearest_region"] = do_info.get("nearestRegion")
            colos[code]["regions"] = do_info.get("regions", {})
        else:
            # Colo exists in DO data but not in components.json
            colos[code] = {
                "name": code,  # Use code as name if unknown
                "cf_region": "Unknown",
                "cf_region_hint": None,
                "nearest_region": do_info.get("nearestRegion"),
                "regions": do_info.get("regions", {}),
            }
    return colos


def sanitize_variant_name(code: str) -> str:
    """Convert colo code to valid Rust enum variant name."""
    # Handle numeric-starting codes like "009" -> "C009"
    if code[0].isdigit():
        return f"C{code}"
    return code


def generate_rust_code(colos: dict[str, dict]) -> str:
    """Generate the Rust source code."""

    # Sort colos by code for consistent output
    sorted_colos = sorted(colos.items())

    lines = [
        "// This file is auto-generated by codegen.py",
        "// Do not edit manually - changes will be overwritten",
        "",
        "use core::fmt;",
        "",
        "/// Durable Objects location hint.",
        "///",
        "/// These are the supported regions for Durable Objects location hints.",
        "/// Note: Location hints are best-effort and not guaranteed.",
        "#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]",
        "#[non_exhaustive]",
        "pub enum LocationHint {",
    ]

    for code, (variant, desc) in sorted(LOCATION_HINTS.items()):
        lines.append(f"    /// {desc}")
        lines.append(f"    {variant},")

    lines.extend([
        "}",
        "",
        "impl LocationHint {",
        "    /// Returns the location hint code as used in the Cloudflare API.",
        "    #[inline]",
        "    pub const fn as_str(&self) -> &'static str {",
        "        match self {",
    ])

    for code, (variant, _) in sorted(LOCATION_HINTS.items()):
        lines.append(f'            Self::{variant} => "{code}",')

    lines.extend([
        "        }",
        "    }",
        "",
        "    /// Returns the human-readable name of the location.",
        "    #[inline]",
        "    pub const fn name(&self) -> &'static str {",
        "        match self {",
    ])

    for code, (variant, desc) in sorted(LOCATION_HINTS.items()):
        lines.append(f'            Self::{variant} => "{desc}",')

    lines.extend([
        "        }",
        "    }",
        "",
        "    /// Returns all location hints.",
        "    pub const ALL: &'static [LocationHint] = &[",
    ])

    for code, (variant, _) in sorted(LOCATION_HINTS.items()):
        lines.append(f"        Self::{variant},")

    lines.extend([
        "    ];",
        "",
        "    /// Parse a location hint from its string code.",
        "    pub fn parse(s: &str) -> Option<Self> {",
        "        match s {",
    ])

    for code, (variant, _) in sorted(LOCATION_HINTS.items()):
        lines.append(f'            "{code}" => Some(Self::{variant}),')

    lines.extend([
        "            _ => None,",
        "        }",
        "    }",
        "}",
        "",
        "impl fmt::Display for LocationHint {",
        "    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {",
        "        f.write_str(self.as_str())",
        "    }",
        "}",
        "",
    ])

    # Generate Colo enum
    lines.extend([
        "/// Cloudflare data center (colo) identifier.",
        "///",
        "/// Each variant represents a Cloudflare edge location.",
        "#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]",
        "#[non_exhaustive]",
        "pub enum Colo {",
    ])

    for code, info in sorted_colos:
        variant = sanitize_variant_name(code)
        name = info.get("name", code)
        # Escape any quotes in the name
        name = name.replace('"', '\\"')
        lines.append(f"    /// {name}")
        lines.append(f"    {variant},")

    lines.extend([
        "}",
        "",
        "impl Colo {",
        "    /// Returns the 3-letter IATA code for this colo.",
        "    #[inline]",
        "    pub const fn code(&self) -> &'static str {",
        "        match self {",
    ])

    for code, _ in sorted_colos:
        variant = sanitize_variant_name(code)
        lines.append(f'            Self::{variant} => "{code}",')

    lines.extend([
        "        }",
        "    }",
        "",
        "    /// Returns the human-readable name of this colo.",
        "    #[inline]",
        "    pub const fn name(&self) -> &'static str {",
        "        match self {",
    ])

    for code, info in sorted_colos:
        variant = sanitize_variant_name(code)
        name = info.get("name", code).replace('"', '\\"')
        lines.append(f'            Self::{variant} => "{name}",')

    lines.extend([
        "        }",
        "    }",
        "",
        "    /// Returns the nearest Durable Objects location hint for this colo.",
        "    ///",
        "    /// This is based on measured latency data and represents the best",
        "    /// location hint to use when creating Durable Objects for requests",
        "    /// arriving at this colo.",
        "    ///",
        "    /// Returns `None` if no mapping data is available for this colo.",
        "    #[inline]",
        "    pub const fn location_hint(&self) -> Option<LocationHint> {",
        "        match self {",
    ])

    for code, info in sorted_colos:
        variant = sanitize_variant_name(code)
        nearest = info.get("nearest_region")
        if nearest and nearest in LOCATION_HINTS:
            hint_variant = LOCATION_HINTS[nearest][0]
            lines.append(f"            Self::{variant} => Some(LocationHint::{hint_variant}),")
        else:
            lines.append(f"            Self::{variant} => None,")

    lines.extend([
        "        }",
        "    }",
        "",
        "    /// Parse a colo from its 3-letter code.",
        "    pub fn from_code(code: &str) -> Option<Self> {",
        "        match code {",
    ])

    for code, _ in sorted_colos:
        variant = sanitize_variant_name(code)
        lines.append(f'            "{code}" => Some(Self::{variant}),')

    lines.extend([
        "            _ => None,",
        "        }",
        "    }",
        "",
        "    /// Returns all colos.",
        f"    pub const ALL: &'static [Colo] = &[",
    ])

    for code, _ in sorted_colos:
        variant = sanitize_variant_name(code)
        lines.append(f"        Self::{variant},")

    lines.extend([
        "    ];",
        "}",
        "",
        "impl fmt::Display for Colo {",
        "    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {",
        "        f.write_str(self.code())",
        "    }",
        "}",
        "",
    ])

    return "\n".join(lines)


def main():
    root = Path(__file__).parent

    # Load data
    components_path = root / "components.json"
    do_path = root / "where.durableobjects.live.json"

    if not components_path.exists():
        print(f"Error: {components_path} not found")
        return 1

    if not do_path.exists():
        print(f"Error: {do_path} not found")
        return 1

    print("Loading components.json...")
    components_data = load_json(components_path)

    print("Loading where.durableobjects.live.json...")
    do_data = load_json(do_path)

    print("Extracting colo data...")
    colos = extract_colos_from_components(components_data)
    print(f"  Found {len(colos)} colos from components.json")

    print("Merging DO region data...")
    colos = merge_do_data(colos, do_data)
    print(f"  Total colos after merge: {len(colos)}")

    # Count colos with region mapping
    with_mapping = sum(1 for c in colos.values() if c.get("nearest_region"))
    print(f"  Colos with region mapping: {with_mapping}")

    print("Generating Rust code...")
    rust_code = generate_rust_code(colos)

    output_path = root / "src" / "generated.rs"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w") as f:
        f.write(rust_code)

    print(f"Generated {output_path}")
    print(f"  {len(rust_code)} bytes")
    print(f"  {rust_code.count(chr(10))} lines")

    return 0


if __name__ == "__main__":
    exit(main())