"""Schema derivation helpers."""

from __future__ import annotations

from collections.abc import Callable
from functools import lru_cache
from types import GenericAlias
from typing import cast

from pydantic import BaseModel, ConfigDict, create_model
from pydantic.dataclasses import dataclass
from sqlmodel import SQLModel

from auen._query import resolve_pk_field
from auen._relations import RelationInfo, inspect_relations
from auen.config import NestedWriteConfig, SchemaConfig
from auen.types import ModelFieldDefinition


class _StrictInputModel(BaseModel):
    """Base model for request schemas that rejects unknown fields."""

    model_config = ConfigDict(extra="forbid")


@dataclass(frozen=True)
class SchemaContainer:
    """Expanded schema bundle including nested create/update schemas."""

    create: type[BaseModel]
    read: type[BaseModel]
    update: type[BaseModel]
    nested_create: type[BaseModel]
    nested_update: type[BaseModel]

    def to_config(self) -> SchemaConfig:
        return SchemaConfig(
            create=self.create,
            read=self.read,
            update=self.update,
            nested_create=self.nested_create,
            nested_update=self.nested_update,
        )


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_flat_schemas(
    model: type[SQLModel],
    *,
    include: set[str] | None = None,
    exclude: set[str] | None = None,
    exclude_pk_from_create: bool = True,
    nested_writes: NestedWriteConfig | None = None,
) -> SchemaConfig:
    pk_fields = _get_pk_field_names(model)
    all_fields = model.model_fields
    exclude_set = exclude or set()

    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_set:
            continue

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

        read_fields[name] = (annotation, default)

        if not (exclude_pk_from_create and name in pk_fields):
            create_fields[name] = (annotation, default)

        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)

    # Backward compatibility: explicit nested relation schema fields in flat models.
    if nested_writes:
        for rel_name, relation_cfg in nested_writes.fields.items():
            related = _derive_flat_schemas(relation_cfg.model)
            create_fields[rel_name] = cast(
                ModelFieldDefinition,
                (related.create | None, None),
            )
            update_fields[rel_name] = cast(
                ModelFieldDefinition,
                (related.update | None, None),
            )
            read_fields[rel_name] = cast(
                ModelFieldDefinition,
                (related.read | None, None),
            )

    _create_model = cast(Callable[..., type[BaseModel]], create_model)
    model_name = model.__name__
    read_schema = _create_model(f"{model_name}Read", __base__=BaseModel, **read_fields)
    create_schema = _create_model(
        f"{model_name}Create", __base__=_StrictInputModel, **create_fields
    )
    update_schema = _create_model(
        f"{model_name}Update", __base__=_StrictInputModel, **update_fields
    )
    return SchemaConfig(create=create_schema, read=read_schema, update=update_schema)


def _augment_update_with_optional_pk(
    model: type[SQLModel], update_schema: type[BaseModel]
) -> type[BaseModel]:
    try:
        pk_name, pk_type = resolve_pk_field(model, None)
    except ValueError:
        return update_schema
    if pk_name in update_schema.model_fields:
        return update_schema
    _create_model = cast(Callable[..., type[BaseModel]], create_model)
    return _create_model(
        f"{model.__name__}NestedNodeUpdate",
        __base__=update_schema,
        **{pk_name: (pk_type | None, None)},
    )


def _resolve_relation_map(
    model: type[SQLModel],
    nested_writes: NestedWriteConfig | None,
    *,
    auto_discover_relations: bool,
) -> dict[str, tuple[type[SQLModel], RelationInfo | None]]:
    relation_map: dict[str, tuple[type[SQLModel], RelationInfo | None]] = {}

    if auto_discover_relations:
        for name, rel_info in inspect_relations(model).items():
            relation_map[name] = (rel_info.target_model, rel_info)

    if nested_writes:
        for rel_name, relation_cfg in nested_writes.fields.items():
            rel_info = relation_map.get(rel_name, (relation_cfg.model, None))[1]
            relation_map[rel_name] = (relation_cfg.model, rel_info)

    return relation_map


def _build_many_to_many_item_schema(
    *,
    model: type[SQLModel],
    relation_name: str,
    relation_info: RelationInfo,
    related_schema: type[BaseModel],
    flat_cache: dict[type[SQLModel], SchemaConfig],
) -> type[BaseModel]:
    _create_model = cast(Callable[..., type[BaseModel]], create_model)

    link_schema_type: type[BaseModel] | None = None
    if relation_info.link_model is not None:
        if relation_info.link_model not in flat_cache:
            flat_cache[relation_info.link_model] = _derive_flat_schemas(
                relation_info.link_model
            )
        link_schema_type = flat_cache[relation_info.link_model].update

    item_fields: dict[str, ModelFieldDefinition] = {
        "entity": cast(ModelFieldDefinition, (related_schema, ...))
    }
    if link_schema_type is not None:
        item_fields["link"] = cast(ModelFieldDefinition, (link_schema_type | None, None))
    else:
        item_fields["link"] = cast(ModelFieldDefinition, (dict[str, object] | None, None))

    return _create_model(
        f"{model.__name__}{relation_name.capitalize()}NestedItem",
        __base__=_StrictInputModel,
        **item_fields,
    )


def _build_nested_update_schema(
    *,
    model: type[SQLModel],
    base_update_schema: type[BaseModel],
    auto_discover_relations: bool,
    nested_writes: NestedWriteConfig | None,
    active_path: tuple[type[SQLModel], ...],
    nested_cache: dict[
        tuple[type[SQLModel], tuple[type[SQLModel], ...]],
        type[BaseModel],
    ],
    flat_cache: dict[type[SQLModel], SchemaConfig],
) -> type[BaseModel]:
    cache_key = (model, active_path)
    if cache_key in nested_cache:
        return nested_cache[cache_key]

    relation_map = _resolve_relation_map(
        model,
        nested_writes,
        auto_discover_relations=auto_discover_relations,
    )

    nested_fields: dict[str, ModelFieldDefinition] = {}
    for relation_name, (related_model, rel_info) in relation_map.items():
        if related_model not in flat_cache:
            flat_cache[related_model] = _derive_flat_schemas(related_model)
        related_base_update = _augment_update_with_optional_pk(
            related_model, flat_cache[related_model].update
        )
        if related_model in active_path:
            related_nested_schema = related_base_update
        else:
            related_nested_schema = _build_nested_update_schema(
                model=related_model,
                base_update_schema=related_base_update,
                auto_discover_relations=auto_discover_relations,
                nested_writes=None,
                active_path=(*active_path, model),
                nested_cache=nested_cache,
                flat_cache=flat_cache,
            )

        if rel_info is not None and rel_info.is_many_to_many:
            item_schema = _build_many_to_many_item_schema(
                model=model,
                relation_name=relation_name,
                relation_info=rel_info,
                related_schema=related_nested_schema,
                flat_cache=flat_cache,
            )
            nested_fields[relation_name] = cast(
                ModelFieldDefinition,
                (GenericAlias(list, item_schema) | None, None),
            )
        elif rel_info is not None and rel_info.is_to_many:
            nested_fields[relation_name] = cast(
                ModelFieldDefinition,
                (GenericAlias(list, related_nested_schema) | None, None),
            )
        else:
            nested_fields[relation_name] = cast(
                ModelFieldDefinition,
                (related_nested_schema | None, None),
            )

    if not nested_fields:
        nested_cache[cache_key] = base_update_schema
        return base_update_schema

    _create_model = cast(Callable[..., type[BaseModel]], create_model)
    nested_schema = _create_model(
        f"{model.__name__}NestedUpdate",
        __base__=base_update_schema,
        **nested_fields,
    )
    nested_cache[cache_key] = nested_schema
    return nested_schema


def _build_nested_create_schema(
    *,
    model: type[SQLModel],
    base_schema: type[BaseModel],
    auto_discover_relations: bool,
    nested_writes: NestedWriteConfig | None,
    active_path: tuple[type[SQLModel], ...],
    nested_cache: dict[
        tuple[type[SQLModel], tuple[type[SQLModel], ...]],
        type[BaseModel],
    ],
    flat_cache: dict[type[SQLModel], SchemaConfig],
) -> type[BaseModel]:
    cache_key = (model, active_path)
    if cache_key in nested_cache:
        return nested_cache[cache_key]

    relation_map = _resolve_relation_map(
        model,
        nested_writes,
        auto_discover_relations=auto_discover_relations,
    )

    nested_fields: dict[str, ModelFieldDefinition] = {}
    for relation_name, (related_model, rel_info) in relation_map.items():
        if related_model not in flat_cache:
            flat_cache[related_model] = _derive_flat_schemas(related_model)
        related_base_schema = _augment_update_with_optional_pk(
            related_model, flat_cache[related_model].update
        )
        if related_model in active_path:
            related_nested_schema = related_base_schema
        else:
            related_nested_schema = _build_nested_create_schema(
                model=related_model,
                base_schema=related_base_schema,
                auto_discover_relations=auto_discover_relations,
                nested_writes=None,
                active_path=(*active_path, model),
                nested_cache=nested_cache,
                flat_cache=flat_cache,
            )

        if rel_info is not None and rel_info.is_many_to_many:
            item_schema = _build_many_to_many_item_schema(
                model=model,
                relation_name=relation_name,
                relation_info=rel_info,
                related_schema=related_nested_schema,
                flat_cache=flat_cache,
            )
            nested_fields[relation_name] = cast(
                ModelFieldDefinition,
                (GenericAlias(list, item_schema) | None, None),
            )
        elif rel_info is not None and rel_info.is_to_many:
            nested_fields[relation_name] = cast(
                ModelFieldDefinition,
                (GenericAlias(list, related_nested_schema) | None, None),
            )
        else:
            nested_fields[relation_name] = cast(
                ModelFieldDefinition,
                (related_nested_schema | None, None),
            )

    if not nested_fields:
        nested_cache[cache_key] = base_schema
        return base_schema

    _create_model = cast(Callable[..., type[BaseModel]], create_model)
    nested_schema = _create_model(
        f"{model.__name__}NestedCreate",
        __base__=base_schema,
        **nested_fields,
    )
    nested_cache[cache_key] = nested_schema
    return nested_schema


def _derive_schema_container_impl(
    model: type[SQLModel],
    *,
    include: set[str] | None,
    exclude: set[str] | None,
    exclude_pk_from_create: bool,
    nested_writes: NestedWriteConfig | None,
    auto_discover_relations: bool,
) -> SchemaContainer:
    flat = _derive_flat_schemas(
        model,
        include=include,
        exclude=exclude,
        exclude_pk_from_create=exclude_pk_from_create,
        nested_writes=nested_writes,
    )
    nested_create = _build_nested_create_schema(
        model=model,
        base_schema=flat.create,
        auto_discover_relations=auto_discover_relations,
        nested_writes=nested_writes,
        active_path=(),
        nested_cache={},
        flat_cache={model: flat},
    )
    nested_update = _build_nested_update_schema(
        model=model,
        base_update_schema=flat.update,
        auto_discover_relations=auto_discover_relations,
        nested_writes=nested_writes,
        active_path=(),
        nested_cache={},
        flat_cache={model: flat},
    )
    return SchemaContainer(
        create=flat.create,
        read=flat.read,
        update=flat.update,
        nested_create=nested_create,
        nested_update=nested_update,
    )


@lru_cache(maxsize=256)
def _derive_schema_container_cached(
    model: type[SQLModel],
    include: frozenset[str] | None,
    exclude: frozenset[str],
    exclude_pk_from_create: bool,
    auto_discover_relations: bool,
) -> SchemaContainer:
    include_set = set(include) if include is not None else None
    exclude_set = set(exclude)
    return _derive_schema_container_impl(
        model,
        include=include_set,
        exclude=exclude_set,
        exclude_pk_from_create=exclude_pk_from_create,
        nested_writes=None,
        auto_discover_relations=auto_discover_relations,
    )


def derive_schema_container(
    model: type[SQLModel],
    *,
    include: set[str] | None = None,
    exclude: set[str] | None = None,
    exclude_pk_from_create: bool = True,
    nested_writes: NestedWriteConfig | None = None,
    auto_discover_relations: bool = True,
) -> SchemaContainer:
    """Derive Create/Read/Update/NestedCreate/NestedUpdate schemas from a SQLModel."""
    include_key = frozenset(include) if include is not None else None
    exclude_key = frozenset(exclude or set())
    if nested_writes is None:
        return _derive_schema_container_cached(
            model,
            include_key,
            exclude_key,
            exclude_pk_from_create,
            auto_discover_relations,
        )
    return _derive_schema_container_impl(
        model,
        include=set(include) if include is not None else None,
        exclude=set(exclude or set()),
        exclude_pk_from_create=exclude_pk_from_create,
        nested_writes=nested_writes,
        auto_discover_relations=auto_discover_relations,
    )


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