"""Python source code generation for config schemas.

Serializes a ``LayerConfigResult`` into a valid, self-contained Python
source file string containing Pydantic ``BaseModel`` class definitions
and an optional discriminated union type alias.

See ``config-typing-design.md`` Section 6.1 for specification.
"""
from __future__ import annotations

from typing import Annotated, Any, Literal, Union, get_args, get_origin

from pydantic import BaseModel
from pydantic.fields import FieldInfo
from pydantic_core import PydanticUndefined

from conscribe.config.builder import LayerConfigResult

# Constraint types from annotated_types (used by Pydantic's Annotated metadata)
try:
    from annotated_types import Ge, Gt, Le, Lt
    _CONSTRAINT_MAP = {Ge: "ge", Gt: "gt", Le: "le", Lt: "lt"}
except ImportError:
    _CONSTRAINT_MAP = {}


def generate_layer_config_source(result: LayerConfigResult) -> str:
    """Serialize a ``LayerConfigResult`` to a valid Python source file string.

    Output structure:
    1. Auto-generated header comment
    2. ``from __future__ import annotations``
    3. Required imports (typing, pydantic)
    4. Per-key class definitions
    5. Union type alias (if multiple models)

    Args:
        result: The ``LayerConfigResult`` from ``build_layer_config()``.

    Returns:
        A string containing valid Python source code.
    """
    parts: list[str] = []

    # 1. Header
    parts.append(_generate_header(result))

    # 2. Future annotations
    # Note: we include this for style consistency but Pydantic needs
    # runtime type evaluation, so we also call model_rebuild() below.
    parts.append("from __future__ import annotations\n")

    # 3. Imports
    parts.append(_generate_imports(result))

    # 4. Per-key class definitions
    model_names = []
    for key in sorted(result.per_key_models.keys()):
        model = result.per_key_models[key]
        parts.append(_generate_class(model, result.discriminator_field))
        model_names.append(model.__name__)

    # 5. model_rebuild() calls (needed because from __future__ import annotations
    #    defers type evaluation; Pydantic needs explicit rebuild)
    rebuild_lines = [f"\n# Rebuild models for deferred annotation resolution"]
    for name in model_names:
        rebuild_lines.append(f"{name}.model_rebuild()")
    rebuild_lines.append("")
    parts.append("\n".join(rebuild_lines))

    # 6. Union type alias (only for multiple models)
    if len(result.per_key_models) > 1:
        parts.append(_generate_union_alias(result))

    return "\n".join(parts)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _generate_header(result: LayerConfigResult) -> str:
    """Generate the auto-generated header comment."""
    count = len(result.per_key_models)
    keys = ", ".join(sorted(result.per_key_models.keys()))
    return (
        f'# Auto-generated by conscribe. DO NOT EDIT.\n'
        f'# Layer: {result.layer_name}\n'
        f'# Discriminator: {result.discriminator_field}\n'
        f'# Entries ({count}): {keys}\n'
    )


def _generate_imports(result: LayerConfigResult) -> str:
    """Generate the import block."""
    # Collect which typing imports are needed
    typing_imports = {"Literal"}
    pydantic_imports = {"BaseModel", "ConfigDict"}

    # Scan fields to determine needed imports
    has_field = False
    for model in result.per_key_models.values():
        for field_name, field_info in model.model_fields.items():
            tp = field_info.annotation
            _collect_type_imports(tp, typing_imports)
            if _needs_field(field_info):
                has_field = True

    if has_field:
        pydantic_imports.add("Field")

    # Check if Union alias is needed
    if len(result.per_key_models) > 1:
        typing_imports.add("Annotated")
        typing_imports.add("Union")
        pydantic_imports.add("Field")

    lines = []
    lines.append(f"from typing import {', '.join(sorted(typing_imports))}")
    lines.append("")
    lines.append(f"from pydantic import {', '.join(sorted(pydantic_imports))}")
    lines.append("")
    return "\n".join(lines)


def _collect_type_imports(tp: Any, imports: set[str]) -> None:
    """Collect typing imports needed for a type annotation."""
    origin = get_origin(tp)
    if origin is Union:
        args = get_args(tp)
        if len(args) == 2 and type(None) in args:
            imports.add("Optional")
        else:
            imports.add("Union")
        for arg in args:
            if arg is not type(None):
                _collect_type_imports(arg, imports)
    elif origin is Annotated:
        args = get_args(tp)
        if args:
            _collect_type_imports(args[0], imports)
    elif tp is Any:
        imports.add("Any")


def _needs_field(field_info: FieldInfo) -> bool:
    """Check if a field needs Field() in its source representation."""
    if field_info.description is not None:
        return True
    if _extract_constraints(field_info):
        return True
    return False


def _extract_constraints(field_info: FieldInfo) -> dict[str, Any]:
    """Extract constraint kwargs from a FieldInfo (checking metadata)."""
    constraints: dict[str, Any] = {}
    # Check metadata for annotated_types constraint objects
    for meta in field_info.metadata:
        for cls, attr_name in _CONSTRAINT_MAP.items():
            if isinstance(meta, cls):
                constraints[attr_name] = getattr(meta, attr_name)
    return constraints


def _generate_class(model: type[BaseModel], disc_field: str) -> str:
    """Generate source code for a single model class."""
    extra = model.model_config.get("extra", "forbid")
    is_open = extra == "allow"

    lines = []
    lines.append(f"\nclass {model.__name__}(BaseModel):")

    # Docstring for open schema
    if is_open:
        lines.append('    """This implementation accepts extra parameters (**kwargs).')
        lines.append('    Undeclared fields will not cause errors."""')
        lines.append("")

    # model_config
    lines.append(f'    model_config = ConfigDict(extra="{extra}")')
    lines.append("")

    # Fields
    for field_name, field_info in model.model_fields.items():
        field_line = _render_field(field_name, field_info)
        lines.append(f"    {field_line}")

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


def _render_field(name: str, field_info: FieldInfo) -> str:
    """Render a single field definition line."""
    type_str = _type_to_source(field_info.annotation)
    default = field_info.default
    is_required = default is PydanticUndefined

    # Check if we need Field()
    field_kwargs: dict[str, Any] = {}
    if field_info.description is not None:
        field_kwargs["description"] = field_info.description
    field_kwargs.update(_extract_constraints(field_info))

    if field_kwargs:
        # Need Field()
        default_repr = "..." if is_required else _value_to_source(default)
        kwargs_str = ", ".join(
            f"{k}={_value_to_source(v)}" for k, v in field_kwargs.items()
        )
        return f"{name}: {type_str} = Field({default_repr}, {kwargs_str})"
    else:
        # Plain default
        if is_required:
            return f"{name}: {type_str} = ..."
        else:
            return f"{name}: {type_str} = {_value_to_source(default)}"


def _type_to_source(tp: Any) -> str:
    """Convert a type annotation to its source code representation."""
    if tp is None or tp is type(None):
        return "None"

    # Handle Annotated
    origin = get_origin(tp)
    if origin is Annotated:
        args = get_args(tp)
        return _type_to_source(args[0]) if args else repr(tp)

    # Handle Literal
    if origin is Literal:
        args = get_args(tp)
        args_str = ", ".join(_value_to_source(a) for a in args)
        return f"Literal[{args_str}]"

    # Handle Union (including Optional)
    if origin is Union:
        args = get_args(tp)
        if len(args) == 2 and type(None) in args:
            # Optional[T]
            non_none = [a for a in args if a is not type(None)][0]
            return f"Optional[{_type_to_source(non_none)}]"
        else:
            args_str = ", ".join(_type_to_source(a) for a in args)
            return f"Union[{args_str}]"

    # Handle generic types (list, dict, etc.)
    if origin is not None:
        args = get_args(tp)
        origin_name = _get_type_name(origin)
        if args:
            args_str = ", ".join(_type_to_source(a) for a in args)
            return f"{origin_name}[{args_str}]"
        return origin_name

    # Basic types
    return _get_type_name(tp)


def _get_type_name(tp: Any) -> str:
    """Get the simple name for a type."""
    if tp is Any:
        return "Any"
    if hasattr(tp, "__name__"):
        return tp.__name__
    return repr(tp)


def _value_to_source(value: Any) -> str:
    """Convert a Python value to its source code representation."""
    if value is None:
        return "None"
    if value is ...:
        return "..."
    if isinstance(value, bool):
        return repr(value)
    if isinstance(value, (int, float)):
        return repr(value)
    if isinstance(value, str):
        return repr(value)
    if isinstance(value, list):
        return repr(value)
    if isinstance(value, dict):
        return repr(value)
    return repr(value)


def _generate_union_alias(result: LayerConfigResult) -> str:
    """Generate the union type alias for multiple models."""
    layer_name = result.layer_name
    if len(layer_name) <= 3:
        alias_name = f"{layer_name.upper()}Config"
    else:
        alias_name = f"{layer_name.title()}Config"

    model_names = [
        result.per_key_models[k].__name__
        for k in sorted(result.per_key_models.keys())
    ]
    union_members = ", ".join(model_names)

    return (
        f'\n{alias_name} = Annotated[\n'
        f'    Union[{union_members}],\n'
        f'    Field(discriminator="{result.discriminator_field}"),\n'
        f']\n'
    )
