"""Schema derivation helpers."""

from __future__ import annotations

from collections.abc import Callable
from typing import cast

from pydantic import BaseModel, create_model
from sqlmodel import SQLModel

from auen.config import SchemaConfig
from auen.types import ModelFieldDefinition


def _get_pk_field_names(model: type[SQLModel]) -> set[str]:
    """Get primary key field names from a SQLModel."""
    from sqlalchemy import inspect as sa_inspect

    try:
        pk_cols = sa_inspect(model).mapper.primary_key
        return {col.name for col in pk_cols}
    except Exception:
        return {"id"}


def derive_schemas(
    model: type[SQLModel],
    *,
    include: set[str] | None = None,
    exclude: set[str] | None = None,
    exclude_pk_from_create: bool = True,
) -> SchemaConfig:
    """Derive Create/Read/Update schemas from a SQLModel.

    This is opt-in because automatic schema derivation can
    be error-prone. Prefer explicit schemas for production.

    Args:
        model: The SQLModel table class.
        include: If set, only include these fields.
        exclude: Fields to exclude from all schemas.
        exclude_pk_from_create: Remove PK from create schema.
    """
    pk_fields = _get_pk_field_names(model)
    all_fields = model.model_fields
    exclude = exclude or set()

    # Build field definitions for each schema type
    read_fields: dict[str, ModelFieldDefinition] = {}
    create_fields: dict[str, ModelFieldDefinition] = {}
    update_fields: dict[str, ModelFieldDefinition] = {}

    for name, field_info in all_fields.items():
        if include and name not in include:
            continue
        if name in exclude:
            continue

        annotation = field_info.annotation
        default = field_info.default
        if field_info.is_required():
            default = ...

        # Read schema includes all fields
        read_fields[name] = (annotation, default)

        # Create schema excludes PKs by default
        if not (exclude_pk_from_create and name in pk_fields):
            create_fields[name] = (annotation, default)

        # Update schema: all non-PK fields optional with default None
        if name not in pk_fields:
            opt_annotation = annotation
            opt_default = default
            if opt_default is ...:
                opt_annotation = annotation | None
                opt_default = None
            update_fields[name] = (opt_annotation, opt_default)

    model_name = model.__name__

    _create_model = cast(Callable[..., type[BaseModel]], create_model)
    read_schema = _create_model(f"{model_name}Read", __base__=BaseModel, **read_fields)
    create_schema = _create_model(
        f"{model_name}Create", __base__=BaseModel, **create_fields
    )
    update_schema = _create_model(
        f"{model_name}Update", __base__=BaseModel, **update_fields
    )

    return SchemaConfig(create=create_schema, read=read_schema, update=update_schema)
