"""Code generation utilities for creating Rust structs from JSON Schema.

This module provides functions to:
- Convert JSON Schema files into Rust struct definitions
- Map JSON Schema types to Rust types
- Generate serde-compatible field attributes
- Handle nullable and optional fields
- Recursively generate nested object types

File: src/python/src/civic_interconnect/cep/codegen/rust_generated.py

uv run cx codegen-rust

uv run python src/python/src/civic_interconnect/cep/codegen/rust_generated.py
"""

import json
from pathlib import Path
import re
from typing import Any

RUST_HEADER = """\
// This file is auto-generated from JSON Schema.
// Do not edit this file by hand. Regenerate via the CLI.

#![allow(clippy::all)]
#![allow(dead_code)]

use serde::{Deserialize, Serialize};

"""

TYPE_MAP = {
    "string": "String",
    "integer": "i64",
    "number": "f64",
    "boolean": "bool",
}


def _rust_ident(name: str) -> str:
    """Convert a JSON property name (camelCase, kebab, etc.) to snake_case for Rust."""
    s = re.sub(r"(.)([A-Z][a-z0-9]+)", r"\1_\2", name)
    s = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", s)
    s = s.replace("-", "_").replace(" ", "_").replace(".", "_")

    out: list[str] = []
    for ch in s:
        if ch.isalnum() or ch == "_":
            out.append(ch.lower())

    ident = "".join(out)

    if not ident:
        ident = "field"
    if ident[0].isdigit():
        ident = "f_" + ident
    return ident


def _to_pascal_case(name: str) -> str:
    """Convert a name to PascalCase for Rust struct names."""
    # Handle kebab-case and snake_case
    parts = re.split(r"[-_]", name)
    return "".join(word.capitalize() for word in parts if word)


def _to_rust_enum_variant(value: str) -> str:
    """Convert a string value (often a URI) to a Rust enum variant name."""
    if "#" in value:
        candidate = value.split("#", 1)[1]
    else:
        candidate = value.rsplit("/", 1)[1] if "/" in value else value

    cleaned = re.sub(r"[^A-Za-z0-9]+", " ", candidate).strip()
    if not cleaned:
        cleaned = "Unknown"

    parts = cleaned.split()
    variant_parts: list[str] = []
    for part in parts:
        if not part:
            continue
        if part[0].isalpha():
            variant_parts.append(part[0].upper() + part[1:].lower())
        else:
            variant_parts.append(part)

    return "".join(variant_parts) or "Unknown"


def _generate_string_enum(enum_name: str, values: list[str]) -> str:
    """Generate a Rust enum for a set of string values."""
    values = [v for v in values if isinstance(v, str)]
    if not values:
        return ""

    lines: list[str] = []
    lines.append("#[derive(Debug, Clone, Serialize, Deserialize)]")
    lines.append("pub enum " + enum_name + " {")

    seen_variants: set[str] = set()
    for v in values:
        variant_name = _to_rust_enum_variant(v)
        if variant_name in seen_variants:
            continue
        seen_variants.add(variant_name)
        lines.append('    #[serde(rename = "' + v + '")]')
        lines.append("    " + variant_name + ",")

    lines.append("}")
    lines.append("")
    return "\n".join(lines)


def _generate_enum_from_vocabulary(vocab_path: Path, enum_name: str) -> str:
    """Generate a Rust enum from a CEP vocabulary file's terms.termUri values."""
    if not vocab_path.exists():
        return ""

    data = json.loads(vocab_path.read_text(encoding="utf-8"))
    terms = data.get("terms", [])
    values: list[str] = []

    for term in terms:
        if isinstance(term, dict):
            uri = term.get("termUri")
            if isinstance(uri, str):
                values.append(uri)

    if not values:
        return ""

    return _generate_string_enum(enum_name, values)


def _type_from_const(prop_schema: dict[str, Any]) -> str | None:
    """Infer a Rust type from a JSON Schema const, if present."""
    if "const" not in prop_schema:
        return None

    const_val = prop_schema["const"]
    explicit_type = prop_schema.get("type")

    if explicit_type in TYPE_MAP:
        return TYPE_MAP[explicit_type]

    if isinstance(const_val, str):
        return TYPE_MAP["string"]
    if isinstance(const_val, bool):
        return TYPE_MAP["boolean"]
    if isinstance(const_val, int):
        return TYPE_MAP["integer"]
    if isinstance(const_val, float):
        return TYPE_MAP["number"]

    return "serde_json::Value"


def _type_from_enum(prop_schema: dict[str, Any]) -> str | None:
    """Infer a Rust type from a JSON Schema enum, if present."""
    if "enum" in prop_schema and prop_schema.get("type") == "string":
        return "String"
    return None


def _type_from_union(
    t: Any,
    prop_schema: dict[str, Any],
    nested_structs: list[str],
    parent_name: str,
    field_name: str,
) -> str | None:
    """Handle union types like ['string', 'null']."""
    if not isinstance(t, list):
        return None

    non_null = [item for item in t if item != "null"]
    if len(non_null) == 1:
        inner = _rust_type_for_schema(
            {"type": non_null[0]}, nested_structs, parent_name, field_name
        )
        return "Option<" + inner + ">"
    return "serde_json::Value"


def _type_from_simple_type(
    t: Any,
    prop_schema: dict[str, Any],
    nested_structs: list[str],
    parent_name: str,
    field_name: str,
) -> str | None:
    """Handle non-union types, including nested objects."""
    if t in TYPE_MAP:
        return TYPE_MAP[t]

    if t == "array":
        items = prop_schema.get("items", {})
        inner = (
            _rust_type_for_schema(items, nested_structs, parent_name, field_name + "Item")
            if items
            else "serde_json::Value"
        )
        return "Vec<" + inner + ">"

    if t == "object":
        props = prop_schema.get("properties")
        if props and isinstance(props, dict):
            # Generate a nested struct
            struct_name = _to_pascal_case(field_name)
            nested_struct = _generate_nested_struct(struct_name, prop_schema, nested_structs)
            nested_structs.append(nested_struct)
            return struct_name
        # No properties defined, use generic value
        return "serde_json::Value"

    return None


def _rust_type_for_schema(
    prop_schema: dict[str, Any],
    nested_structs: list[str],
    parent_name: str = "",
    field_name: str = "",
) -> str:
    """Map a JSON Schema property to a Rust type string."""
    const_type = _type_from_const(prop_schema)
    if const_type is not None:
        return const_type

    enum_type = _type_from_enum(prop_schema)
    if enum_type is not None:
        return enum_type

    t = prop_schema.get("type")
    union_type = _type_from_union(t, prop_schema, nested_structs, parent_name, field_name)
    if union_type is not None:
        return union_type

    simple_type = _type_from_simple_type(t, prop_schema, nested_structs, parent_name, field_name)
    if simple_type is not None:
        return simple_type

    return "serde_json::Value"


def _generate_nested_struct(
    struct_name: str,
    schema: dict[str, Any],
    nested_structs: list[str],
) -> str:
    """Generate a Rust struct for a nested object schema."""
    properties: dict[str, Any] = schema.get("properties", {})
    required: list[str] = schema.get("required", [])
    description = schema.get("description", "")

    lines: list[str] = []

    if description:
        lines.append("/// " + description)

    lines.append("#[derive(Debug, Clone, Serialize, Deserialize)]")
    lines.append('#[serde(rename_all = "camelCase")]')
    lines.append("pub struct " + struct_name + " {")

    for name, prop_schema in properties.items():
        rust_name = _rust_ident(name)
        ty = _rust_type_for_schema(prop_schema, nested_structs, struct_name, name)

        is_required = name in required
        if not is_required and not ty.startswith("Option<"):
            ty = "Option<" + ty + ">"

        # Add serde rename if needed
        if rust_name != name:
            lines.append('    #[serde(rename = "' + name + '")]')

        lines.append("    pub " + rust_name + ": " + ty + ",")

    lines.append("}")
    lines.append("")
    return "\n".join(lines)


def _resolve_ref(ref: str, root: dict[str, Any]) -> dict[str, Any]:
    """Resolve a local JSON Pointer like '#/definitions/X' into a schema dict."""
    if not ref.startswith("#/"):
        return {}
    pointer = ref[2:]
    parts = pointer.split("/") if pointer else []
    node: Any = root
    for part in parts:
        token = part.replace("~1", "/").replace("~0", "~")
        if isinstance(node, dict) and token in node:
            node = node[token]
        else:
            return {}
    if isinstance(node, dict):
        return node
    return {}


def _merge_schema(target: dict[str, Any], source: dict[str, Any]) -> None:
    """Merge properties and required fields from source into target."""
    src_props = source.get("properties", {})
    if src_props:
        props = target.setdefault("properties", {})
        for name, prop_schema in src_props.items():
            props[name] = prop_schema

    src_required = source.get("required", [])
    if src_required:
        required = target.setdefault("required", [])
        for name in src_required:
            if name not in required:
                required.append(name)


def _flatten_schema(schema: dict[str, Any], root: dict[str, Any], out: dict[str, Any]) -> None:
    """Populate out with a flattened view of properties/required via allOf and local refs."""
    _merge_schema(out, schema)

    all_of = schema.get("allOf", [])
    for sub in all_of:
        if not isinstance(sub, dict):
            continue
        if "$ref" in sub:
            ref_schema = _resolve_ref(sub["$ref"], root)
            if ref_schema:
                _flatten_schema(ref_schema, root, out)
        else:
            _flatten_schema(sub, root, out)


def _is_required(name: str, flat_schema: dict[str, Any]) -> bool:
    required = flat_schema.get("required", [])
    return name in required


def _field_line(
    name: str,
    prop_schema: dict[str, Any],
    flat_schema: dict[str, Any],
    special_types: dict[str, str],
    nested_structs: list[str],
    parent_name: str,
) -> str:
    """Generate a single Rust field line for a given JSON Schema property."""
    rust_name = _rust_ident(name)

    if name in special_types:
        ty = special_types[name]
    else:
        ty = _rust_type_for_schema(prop_schema, nested_structs, parent_name, name)

    is_required = _is_required(name, flat_schema)

    if not is_required and not ty.startswith("Option<"):
        ty = "Option<" + ty + ">"

    attrs: list[str] = []
    if rust_name != name:
        attrs.append('#[serde(rename = "' + name + '")]')

    attr_block = ""
    if attrs:
        attr_block = "\n    " + "\n    ".join(attrs)

    return attr_block + "\n    pub " + rust_name + ": " + ty + ","


def _load_schema_for_codegen(schema_path: Path) -> dict[str, Any]:
    """Load a schema and, if it references the shared record envelope, merge that in."""
    raw: dict[str, Any] = json.loads(schema_path.read_text(encoding="utf-8"))

    all_of = raw.get("allOf", [])
    has_envelope_ref = False
    for item in all_of:
        if isinstance(item, dict):
            ref = item.get("$ref")
            if isinstance(ref, str) and "cep.record-envelope.schema.json" in ref:
                has_envelope_ref = True
                break

    if not has_envelope_ref:
        return raw

    envelope_path = schema_path.parent / "cep.record-envelope.schema.json"
    if not envelope_path.exists():
        return raw

    envelope: dict[str, Any] = json.loads(envelope_path.read_text(encoding="utf-8"))

    combined: dict[str, Any] = {
        "description": raw.get("description") or envelope.get("description", ""),
        "allOf": [envelope, raw],
    }

    if isinstance(envelope.get("$defs"), dict):
        combined["$defs"] = envelope["$defs"]

    return combined


def _generate_simple_struct(
    schema: dict[str, Any],
    root: dict[str, Any],
    struct_name: str,
    special_types: dict[str, str],
    nested_structs: list[str],
) -> str:
    """Generate a Rust struct for a simple object schema (no outer envelope merging)."""
    flat: dict[str, Any] = {}
    _flatten_schema(schema, root, flat)

    properties: dict[str, Any] = flat.get("properties", {})
    description = schema.get("description", "")

    lines: list[str] = []

    if description:
        lines.append("/// " + description + "\n")

    lines.append("#[derive(Debug, Clone, Serialize, Deserialize)]")
    lines.append('#[serde(rename_all = "camelCase")]')
    lines.append("pub struct " + struct_name + " {")

    for name, prop_schema in properties.items():
        lines.append(
            _field_line(name, prop_schema, flat, special_types, nested_structs, struct_name)
        )

    lines.append("}")
    lines.append("")
    return "\n".join(lines)


def _generate_record_kind_enum(root: dict[str, Any]) -> str:
    """Generate RecordKind enum from envelope's recordKind enum, if present."""
    all_of = root.get("allOf", [])
    for sub in all_of:
        if not isinstance(sub, dict):
            continue
        props = sub.get("properties", {})
        record_kind_schema = props.get("recordKind")
        if not isinstance(record_kind_schema, dict):
            continue
        values = record_kind_schema.get("enum")
        if isinstance(values, list):
            return _generate_string_enum(
                "RecordKind",
                [v for v in values if isinstance(v, str)],
            )
    return ""


def _generate_status_code_enum_from_defs(
    defs: dict[str, Any],
    special_types: dict[str, str],
) -> str:
    """Generate StatusCode enum from $defs.status.properties.statusCode.enum."""
    status_schema = defs.get("status")
    status_code_enum_values: list[str] = []

    if isinstance(status_schema, dict):
        status_props = status_schema.get("properties", {})
        status_code_schema = status_props.get("statusCode")
        if isinstance(status_code_schema, dict):
            enum_vals = status_code_schema.get("enum")
            if isinstance(enum_vals, list):
                status_code_enum_values = [v for v in enum_vals if isinstance(v, str)]

    if not status_code_enum_values:
        return ""

    enum_code = _generate_string_enum("StatusCode", status_code_enum_values)
    special_types["statusCode"] = "StatusCode"
    return enum_code


def _generate_envelope_structs_from_defs(
    defs: dict[str, Any],
    root: dict[str, Any],
    special_types: dict[str, str],
    nested_structs: list[str],
) -> str:
    """Generate StatusEnvelope, Timestamps, Attestation from $defs."""
    pieces: list[str] = []

    mapping = [
        ("status", "StatusEnvelope"),
        ("timestamps", "Timestamps"),
        ("attestation", "Attestation"),
    ]

    for key, struct_name in mapping:
        schema = defs.get(key)
        if isinstance(schema, dict):
            pieces.append(
                _generate_simple_struct(schema, root, struct_name, special_types, nested_structs)
            )

    special_types["status"] = "StatusEnvelope"
    special_types["timestamps"] = "Timestamps"
    special_types["attestations"] = "Vec<Attestation>"
    special_types["attestation"] = "Attestation"

    return "\n".join(pieces)


def _generate_envelope_items(
    root: dict[str, Any], special_types: dict[str, str], nested_structs: list[str]
) -> str:
    """Generate enums/structs from the shared envelope $defs if present."""
    defs = root.get("$defs")
    if not isinstance(defs, dict):
        return ""

    pieces: list[str] = []

    status_code_enum = _generate_status_code_enum_from_defs(defs, special_types)
    if status_code_enum:
        pieces.append(status_code_enum)

    env_structs = _generate_envelope_structs_from_defs(defs, root, special_types, nested_structs)
    if env_structs:
        pieces.append(env_structs)

    return "\n".join(pieces)


def generate_rust_struct(schema_path: Path, struct_name: str) -> str:
    """Generate a single Rust struct from a JSON Schema file."""
    root: dict[str, Any] = _load_schema_for_codegen(schema_path)

    flat: dict[str, Any] = {}
    _flatten_schema(root, root, flat)

    properties: dict[str, Any] = flat.get("properties", {})
    description = root.get("description", "")

    lines: list[str] = [RUST_HEADER]

    special_types: dict[str, str] = {}
    nested_structs: list[str] = []

    record_kind_enum = _generate_record_kind_enum(root)
    if record_kind_enum:
        lines.append(record_kind_enum)
        special_types["recordKind"] = "RecordKind"

    envelope_items = _generate_envelope_items(root, special_types, nested_structs)
    if envelope_items:
        lines.append(envelope_items)

    identifiers_items = ""
    if "identifiers" in properties:
        identifiers_items = (
            "#[derive(Debug, Clone, Serialize, Deserialize)]\n"
            '#[serde(rename_all = "camelCase")]\n'
            "pub struct Identifier {\n"
            '    #[serde(rename = "schemeUri")]\n'
            "    pub scheme_uri: String,\n"
            "    pub identifier: String,\n"
            '    #[serde(rename = "sourceReference", skip_serializing_if = "Option::is_none")]\n'
            "    pub source_reference: Option<String>,\n"
            "}\n"
            "\n"
            "pub type Identifiers = Vec<Identifier>;\n"
            "\n"
        )
        special_types["identifiers"] = "Identifiers"

    if identifiers_items:
        lines.append(identifiers_items)

    # Process all fields to collect nested structs
    for name, prop_schema in properties.items():
        if name not in special_types:
            _rust_type_for_schema(prop_schema, nested_structs, struct_name, name)

    # Emit nested structs before the main struct
    for nested in nested_structs:
        lines.append(nested)

    # Emit the main record struct
    if description:
        lines.append("/// " + description + "\n")

    lines.append("#[derive(Debug, Clone, Serialize, Deserialize)]")
    lines.append('#[serde(rename_all = "camelCase")]')
    lines.append("pub struct " + struct_name + " {")

    for name, prop_schema in properties.items():
        lines.append(
            _field_line(name, prop_schema, flat, special_types, nested_structs, struct_name)
        )

    lines.append("}")
    lines.append("")
    return "\n".join(lines)


def write_generated_rust(
    schema_path: Path,
    struct_name: str,
    out_path: Path,
) -> None:
    """Generate Rust code from schema_path and write to out_path."""
    out_path.parent.mkdir(parents=True, exist_ok=True)
    rust_code = generate_rust_struct(schema_path, struct_name)
    out_path.write_text(rust_code, encoding="utf-8")
