"""Shared helpers used across all code generators."""

from __future__ import annotations

from ..ir import (
    WitEnum,
    WitList,
    WitOption,
    WitPrimitive,
    WitRecord,
    WitRef,
    WitResult,
    WitType,
    WitTypeAlias,
    WitTypeDef,
    WitVariant,
    WitVariantCase,
    WitWorld,
)
from ..naming import kebab_to_pascal

GENERATED_HEADER = '''\
"""GENERATED by witgen — DO NOT EDIT.

Re-generate with: uv run python -m tools.witgen
"""
'''


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _is_unit_variant_case(case: WitVariantCase) -> bool:
    return case.payload is None


def _is_named_record_ref(ty: WitType) -> bool:
    """True if the type is a reference to a named record."""
    return isinstance(ty, WitRef)


def _variant_has_overlapping_record_cases(
    variant: WitVariant, type_map: dict[str, WitTypeDef]
) -> bool:
    """Check if all variant cases reference distinct named records.

    Returns False if any case lacks a payload, isn't a ref to a record,
    or if two cases reference the same record type (which would make the
    Union representation ambiguous for to_wasm_* dispatch).
    """
    seen_names: set[str] = set()
    for case in variant.cases:
        if case.payload is None:
            return False
        if not isinstance(case.payload, WitRef):
            return False
        referred = type_map.get(case.payload.name)
        if not isinstance(referred, WitRecord):
            return False
        if case.payload.name in seen_names:
            return False
        seen_names.add(case.payload.name)
    return True


def _resolve_type(ty: WitType, type_map: dict[str, WitTypeDef]) -> WitType:
    """Resolve type aliases."""
    if isinstance(ty, WitRef):
        defn = type_map.get(ty.name)
        if isinstance(defn, WitTypeAlias):
            return _resolve_type(defn.target, type_map)
    return ty


def _build_type_map(world: WitWorld) -> dict[str, WitTypeDef]:
    """Build name->definition map from world types + imported types."""
    m: dict[str, WitTypeDef] = {}
    for t in world.types:
        m[t.name] = t
    for imp in world.imports:
        for t in imp.types:
            m[t.name] = t
    return m


def _py_type_annotation(
    ty: WitType,
    type_map: dict[str, WitTypeDef],
) -> str:
    """Generate a Python type annotation for a WIT type."""
    ty = _resolve_type(ty, type_map)
    if isinstance(ty, WitPrimitive):
        if ty.name == "string":
            return "str"
        if ty.name == "bool":
            return "bool"
        if ty.name in ("u32", "u64", "s32", "s64"):
            return "int"
        if ty.name in ("f32", "f64"):
            return "float"
        return ty.name
    if isinstance(ty, WitList):
        inner = _py_type_annotation(ty.element, type_map)
        return f"list[{inner}]"
    if isinstance(ty, WitOption):
        inner = _py_type_annotation(ty.inner, type_map)
        return f"{inner} | None"
    if isinstance(ty, WitResult):
        ok_ann = _py_type_annotation(ty.ok, type_map) if ty.ok else "None"
        err_ann = _py_type_annotation(ty.err, type_map) if ty.err else "None"
        return f"Result[{ok_ann}, {err_ann}]"
    if isinstance(ty, WitRef):
        defn = type_map.get(ty.name)
        if isinstance(defn, WitEnum):
            return "str"
        if isinstance(defn, WitVariant):
            return kebab_to_pascal(ty.name)
        if isinstance(defn, WitRecord):
            return kebab_to_pascal(ty.name)
        return kebab_to_pascal(ty.name)
    return "Any"


# Module-level field overrides, set once before generation via init_field_overrides().
_field_overrides: dict[str, str] = {}


def init_field_overrides(overrides: dict[str, str]) -> None:
    """Set variant case field name overrides from config. Call once before generating."""
    _field_overrides.clear()
    _field_overrides.update(overrides)


def _variant_payload_field_name(
    variant_name: str,
    case_name: str,
    payload: WitType,
    type_map: dict[str, WitTypeDef],
) -> str:
    """Determine the field name for a variant case payload.

    Returns the name from ``field_overrides`` if configured for this
    ``variant_name.case_name`` key, otherwise defaults to ``'value'``.
    """
    key = f"{variant_name}.{case_name}"
    if key in _field_overrides:
        return _field_overrides[key]
    return "value"


def _record_has_variant_fields(
    record: WitRecord, type_map: dict[str, WitTypeDef]
) -> bool:
    for field in record.fields:
        resolved = _resolve_type(field.type, type_map)
        if isinstance(resolved, WitRef):
            defn = type_map.get(resolved.name)
            if isinstance(defn, WitVariant):
                return True
    return False


def _record_needs_from_wasm(record: WitRecord, type_map: dict[str, WitTypeDef]) -> bool:
    """Check if a record needs a from_wasm converter (has multi-word fields or variant fields)."""
    for field in record.fields:
        if "-" in field.name:
            return True
        if isinstance(field.type, WitRef):
            defn = type_map.get(field.type.name)
            if isinstance(defn, WitVariant):
                return True
            if isinstance(defn, WitRecord) and _record_needs_from_wasm(defn, type_map):
                return True
        elif isinstance(field.type, WitList):
            if isinstance(field.type.element, WitRef):
                inner_defn = type_map.get(field.type.element.name)
                if isinstance(inner_defn, WitRecord) and _record_needs_from_wasm(
                    inner_defn, type_map
                ):
                    return True
    return False


def _collect_type_refs(
    ty: WitType,
    type_map: dict[str, WitTypeDef],
    names: set[str],
) -> None:
    """Recursively collect Python type names referenced by a WIT type."""
    if isinstance(ty, WitRef):
        defn = type_map.get(ty.name)
        if isinstance(defn, WitEnum):
            return  # enums are str
        if isinstance(defn, WitTypeAlias):
            _collect_type_refs(defn.target, type_map, names)
            return
        if isinstance(defn, WitVariant):
            if _variant_has_overlapping_record_cases(defn, type_map):
                names.add(kebab_to_pascal(ty.name))
                for case in defn.cases:
                    assert isinstance(case.payload, WitRef)
                    names.add(kebab_to_pascal(case.payload.name))
            else:
                names.add(kebab_to_pascal(ty.name))
        elif isinstance(defn, WitRecord):
            names.add(kebab_to_pascal(ty.name))
    elif isinstance(ty, WitResult):
        names.add("Result")
        if ty.ok:
            _collect_type_refs(ty.ok, type_map, names)
        if ty.err:
            _collect_type_refs(ty.err, type_map, names)
    elif isinstance(ty, WitList):
        _collect_type_refs(ty.element, type_map, names)
    elif isinstance(ty, WitOption):
        _collect_type_refs(ty.inner, type_map, names)


def _collect_annotation_type_refs(
    ty: WitType,
    type_map: dict[str, WitTypeDef],
    names: set[str],
) -> None:
    """Collect Python type names that appear in annotations (no variant member expansion)."""
    if isinstance(ty, WitRef):
        defn = type_map.get(ty.name)
        if isinstance(defn, WitEnum):
            return  # enums are str
        if isinstance(defn, WitTypeAlias):
            _collect_annotation_type_refs(defn.target, type_map, names)
            return
        names.add(kebab_to_pascal(ty.name))
    elif isinstance(ty, WitResult):
        names.add("Result")
        if ty.ok:
            _collect_annotation_type_refs(ty.ok, type_map, names)
        if ty.err:
            _collect_annotation_type_refs(ty.err, type_map, names)
    elif isinstance(ty, WitList):
        _collect_annotation_type_refs(ty.element, type_map, names)
    elif isinstance(ty, WitOption):
        _collect_annotation_type_refs(ty.inner, type_map, names)
