"""Nested mutation execution for /nested endpoints."""

from __future__ import annotations

from dataclasses import dataclass
from typing import NoReturn, cast

from fastapi import HTTPException
from pydantic import BaseModel, ValidationError
from sqlmodel import SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession

from auen._query import resolve_pk_field
from auen._relations import RelationInfo, inspect_relations
from auen.config import NestedWriteConfig
from auen.exceptions import ConfigurationError, ForbiddenError, NotFoundError
from auen.policy import AllowAll, CrudPolicy
from auen.schemas import derive_schema_container
from auen.types import PKValue, User


@dataclass
class _MutationContext:
    session: AsyncSession
    user: User
    nested_policy_registry: dict[type[SQLModel], CrudPolicy]
    relation_policy_overrides: dict[tuple[type[SQLModel], str], CrudPolicy]


def _unprocessable(detail: str) -> NoReturn:
    raise HTTPException(status_code=422, detail=detail)


def _as_dict(value: object, *, relation_name: str) -> dict[str, object]:
    if not isinstance(value, dict):
        _unprocessable(f"Field '{relation_name}' must be an object for nested writes")
    return cast(dict[str, object], value)


def _as_list(value: object, *, relation_name: str) -> list[object]:
    if not isinstance(value, list):
        _unprocessable(f"Field '{relation_name}' must be a list for nested writes")
    return cast(list[object], value)


def _policy_for_model(
    ctx: _MutationContext, model: type[SQLModel], *, fallback: CrudPolicy | None = None
) -> CrudPolicy:
    if fallback is not None:
        return fallback
    return ctx.nested_policy_registry.get(model, AllowAll())


def _policy_for_relation(
    ctx: _MutationContext,
    parent_model: type[SQLModel],
    relation_name: str,
    target_model: type[SQLModel],
) -> CrudPolicy:
    override = ctx.relation_policy_overrides.get((parent_model, relation_name))
    return _policy_for_model(ctx, target_model, fallback=override)


def _flat_schemas(model: type[SQLModel]) -> tuple[type[BaseModel], type[BaseModel]]:
    container = derive_schema_container(model, auto_discover_relations=False)
    return container.create, container.update


def _validate_model(schema: type[BaseModel], payload: dict[str, object]) -> BaseModel:
    try:
        return schema.model_validate(payload)
    except ValidationError as exc:
        _unprocessable(str(exc))


def _split_payload(
    model: type[SQLModel],
    payload: dict[str, object],
) -> tuple[dict[str, object], dict[str, object], dict[str, RelationInfo]]:
    relations = inspect_relations(model)
    scalar_payload: dict[str, object] = {}
    relation_payload: dict[str, object] = {}
    for key, value in payload.items():
        if key in relations:
            relation_payload[key] = value
        else:
            scalar_payload[key] = value
    return scalar_payload, relation_payload, relations


def _many_to_one_fk_field(rel_info: RelationInfo) -> str:
    if len(rel_info.local_columns) != 1:
        msg = (
            f"Nested writes require exactly one local FK column for "
            f"{rel_info.parent_model.__name__}.{rel_info.name}"
        )
        raise ConfigurationError(msg)
    return rel_info.local_columns[0]


def _one_to_many_pair(rel_info: RelationInfo) -> tuple[str, str]:
    if len(rel_info.sync_pairs) != 1:
        msg = (
            f"Nested writes require exactly one sync pair for "
            f"{rel_info.parent_model.__name__}.{rel_info.name}"
        )
        raise ConfigurationError(msg)
    return rel_info.sync_pairs[0]


def _many_to_many_pairs(
    rel_info: RelationInfo,
) -> tuple[tuple[str, str], tuple[str, str]]:
    if len(rel_info.sync_pairs) != 1 or len(rel_info.secondary_sync_pairs) != 1:
        msg = (
            f"Nested writes require exactly one parent and one target sync pair for "
            f"{rel_info.parent_model.__name__}.{rel_info.name}"
        )
        raise ConfigurationError(msg)
    return rel_info.sync_pairs[0], rel_info.secondary_sync_pairs[0]


async def _create_model_graph(
    ctx: _MutationContext,
    model: type[SQLModel],
    payload: dict[str, object],
    *,
    policy_override: CrudPolicy | None = None,
) -> SQLModel:
    scalar_data, relation_data, relations = _split_payload(model, payload)
    create_schema, _update_schema = _flat_schemas(model)

    # Resolve many-to-one relations before creating this object so FK values are present.
    for relation_name, relation_payload in list(relation_data.items()):
        rel_info = relations[relation_name]
        if not rel_info.is_many_to_one:
            continue
        await _apply_many_to_one(
            ctx,
            parent_model=model,
            relation_name=relation_name,
            rel_info=rel_info,
            relation_payload=relation_payload,
            parent_obj=None,
            parent_scalar_data=scalar_data,
        )
        relation_data.pop(relation_name)

    create_obj = _validate_model(create_schema, scalar_data)
    policy = _policy_for_model(ctx, model, fallback=policy_override)
    if not policy.can_create(ctx.user, create_obj):
        raise ForbiddenError()

    db_obj = model(**create_obj.model_dump())
    ctx.session.add(db_obj)
    await ctx.session.flush()

    for relation_name, relation_payload in relation_data.items():
        rel_info = relations[relation_name]
        await _apply_relation(
            ctx,
            parent_model=model,
            parent_obj=db_obj,
            relation_name=relation_name,
            rel_info=rel_info,
            relation_payload=relation_payload,
            parent_scalar_data=scalar_data,
        )

    return db_obj


async def _update_model_graph(
    ctx: _MutationContext,
    model: type[SQLModel],
    db_obj: SQLModel,
    payload: dict[str, object],
    *,
    policy_override: CrudPolicy | None = None,
) -> None:
    scalar_data, relation_data, relations = _split_payload(model, payload)
    _create_schema, update_schema = _flat_schemas(model)

    for relation_name, relation_payload in relation_data.items():
        rel_info = relations[relation_name]
        await _apply_relation(
            ctx,
            parent_model=model,
            parent_obj=db_obj,
            relation_name=relation_name,
            rel_info=rel_info,
            relation_payload=relation_payload,
            parent_scalar_data=scalar_data,
        )

    update_obj = _validate_model(update_schema, scalar_data)
    policy = _policy_for_model(ctx, model, fallback=policy_override)
    if not policy.can_update(ctx.user, db_obj, update_obj):
        raise ForbiddenError()

    for field, value in update_obj.model_dump(exclude_unset=True).items():
        setattr(db_obj, field, value)
    ctx.session.add(db_obj)
    await ctx.session.flush()


async def _apply_relation(
    ctx: _MutationContext,
    *,
    parent_model: type[SQLModel],
    parent_obj: SQLModel | None,
    relation_name: str,
    rel_info: RelationInfo,
    relation_payload: object,
    parent_scalar_data: dict[str, object],
) -> None:
    if rel_info.is_many_to_one:
        await _apply_many_to_one(
            ctx,
            parent_model=parent_model,
            relation_name=relation_name,
            rel_info=rel_info,
            relation_payload=relation_payload,
            parent_obj=parent_obj,
            parent_scalar_data=parent_scalar_data,
        )
    elif rel_info.is_many_to_many:
        if parent_obj is None:
            _unprocessable(f"Field '{relation_name}' requires a persisted parent object")
        await _apply_many_to_many(
            ctx,
            parent_model=parent_model,
            parent_obj=parent_obj,
            relation_name=relation_name,
            rel_info=rel_info,
            relation_payload=relation_payload,
        )
    elif rel_info.is_one_to_many and rel_info.uselist:
        if parent_obj is None:
            _unprocessable(f"Field '{relation_name}' requires a persisted parent object")
        await _apply_one_to_many(
            ctx,
            parent_model=parent_model,
            parent_obj=parent_obj,
            relation_name=relation_name,
            rel_info=rel_info,
            relation_payload=relation_payload,
        )
    elif rel_info.is_one_to_many and not rel_info.uselist:
        if parent_obj is None:
            _unprocessable(f"Field '{relation_name}' requires a persisted parent object")
        payload = _as_dict(relation_payload, relation_name=relation_name)
        await _apply_one_to_many(
            ctx,
            parent_model=parent_model,
            parent_obj=parent_obj,
            relation_name=relation_name,
            rel_info=rel_info,
            relation_payload=[payload],
        )
    else:
        _unprocessable(f"Unsupported nested relation type for '{relation_name}'")


async def _apply_many_to_one(
    ctx: _MutationContext,
    *,
    parent_model: type[SQLModel],
    relation_name: str,
    rel_info: RelationInfo,
    relation_payload: object,
    parent_obj: SQLModel | None,
    parent_scalar_data: dict[str, object],
) -> None:
    nested_payload = _as_dict(relation_payload, relation_name=relation_name)
    fk_field = _many_to_one_fk_field(rel_info)
    target_model = rel_info.target_model
    target_policy = _policy_for_relation(ctx, parent_model, relation_name, target_model)

    related_pk_name, _related_pk_type = resolve_pk_field(target_model, None)
    related_pk = nested_payload.get(related_pk_name)
    if fk_field in parent_scalar_data:
        if related_pk is None:
            _unprocessable(
                f"Cannot send '{fk_field}' with '{relation_name}' without nested PK"
            )
        if parent_scalar_data[fk_field] != related_pk:
            _unprocessable(
                f"Conflict between '{fk_field}' and '{relation_name}.{related_pk_name}'"
            )

    payload_without_pk = {
        key: value for key, value in nested_payload.items() if key != related_pk_name
    }
    if related_pk is None:
        related_obj = await _create_model_graph(
            ctx,
            target_model,
            payload_without_pk,
            policy_override=target_policy,
        )
    else:
        related_pk_value = cast(PKValue, related_pk)
        related_obj = await ctx.session.get(target_model, related_pk_value)
        if related_obj is None:
            raise NotFoundError(
                target_model.__name__,
                related_pk_name,
                related_pk_value,
            )
        await _update_model_graph(
            ctx,
            target_model,
            related_obj,
            payload_without_pk,
            policy_override=target_policy,
        )

    related_pk_value = getattr(related_obj, related_pk_name)
    parent_scalar_data[fk_field] = related_pk_value
    if parent_obj is not None:
        setattr(parent_obj, fk_field, related_pk_value)


async def _apply_one_to_many(
    ctx: _MutationContext,
    *,
    parent_model: type[SQLModel],
    parent_obj: SQLModel,
    relation_name: str,
    rel_info: RelationInfo,
    relation_payload: object,
) -> None:
    items = _as_list(relation_payload, relation_name=relation_name)
    parent_pk_col, child_fk_col = _one_to_many_pair(rel_info)
    parent_pk_value = getattr(parent_obj, parent_pk_col)
    target_model = rel_info.target_model
    target_policy = _policy_for_relation(ctx, parent_model, relation_name, target_model)
    target_pk_name, _target_pk_type = resolve_pk_field(target_model, None)

    if parent_pk_value is None:
        await ctx.session.flush()
        parent_pk_value = getattr(parent_obj, parent_pk_col)

    for item in items:
        nested_payload = _as_dict(item, relation_name=relation_name)
        target_pk = nested_payload.get(target_pk_name)
        payload = {k: v for k, v in nested_payload.items() if k != target_pk_name}
        if child_fk_col in payload and payload[child_fk_col] != parent_pk_value:
            _unprocessable(
                f"Conflict between '{child_fk_col}' and parent '{parent_pk_col}'"
            )
        payload[child_fk_col] = parent_pk_value

        if target_pk is None:
            await _create_model_graph(
                ctx,
                target_model,
                payload,
                policy_override=target_policy,
            )
            continue

        target_pk_value = cast(PKValue, target_pk)
        child_obj = await ctx.session.get(target_model, target_pk_value)
        if child_obj is None:
            raise NotFoundError(
                target_model.__name__,
                target_pk_name,
                target_pk_value,
            )
        existing_fk = getattr(child_obj, child_fk_col)
        if existing_fk is not None and existing_fk != parent_pk_value:
            _unprocessable(
                f"{target_model.__name__}.{target_pk_name}={target_pk!r} is not related "
                f"to {parent_model.__name__}.{parent_pk_col}={parent_pk_value!r}"
            )
        await _update_model_graph(
            ctx,
            target_model,
            child_obj,
            payload,
            policy_override=target_policy,
        )


async def _apply_many_to_many(
    ctx: _MutationContext,
    *,
    parent_model: type[SQLModel],
    parent_obj: SQLModel,
    relation_name: str,
    rel_info: RelationInfo,
    relation_payload: object,
) -> None:
    items = _as_list(relation_payload, relation_name=relation_name)
    (parent_pk_col, link_parent_fk), (target_pk_col, link_target_fk) = (
        _many_to_many_pairs(rel_info)
    )
    parent_pk_value = getattr(parent_obj, parent_pk_col)
    target_model = rel_info.target_model
    target_policy = _policy_for_relation(ctx, parent_model, relation_name, target_model)
    target_pk_name, _target_pk_type = resolve_pk_field(target_model, None)
    link_model = rel_info.link_model
    link_policy = (
        _policy_for_model(ctx, link_model)
        if link_model is not None
        else cast(CrudPolicy, AllowAll())
    )

    if parent_pk_value is None:
        await ctx.session.flush()
        parent_pk_value = getattr(parent_obj, parent_pk_col)

    for item in items:
        envelope = _as_dict(item, relation_name=relation_name)
        if "entity" not in envelope:
            _unprocessable(f"Field '{relation_name}' items must contain 'entity'")
        entity_payload = _as_dict(
            envelope["entity"],
            relation_name=f"{relation_name}.entity",
        )
        link_payload_raw = envelope.get("link")
        link_payload = (
            _as_dict(link_payload_raw, relation_name=f"{relation_name}.link")
            if link_payload_raw is not None
            else {}
        )

        target_pk = entity_payload.get(target_pk_name)
        entity_data = {
            key: value for key, value in entity_payload.items() if key != target_pk_name
        }
        if target_pk is None:
            target_obj = await _create_model_graph(
                ctx,
                target_model,
                entity_data,
                policy_override=target_policy,
            )
        else:
            target_pk_value = cast(PKValue, target_pk)
            target_obj = await ctx.session.get(target_model, target_pk_value)
            if target_obj is None:
                raise NotFoundError(
                    target_model.__name__,
                    target_pk_name,
                    target_pk_value,
                )
            await _update_model_graph(
                ctx,
                target_model,
                target_obj,
                entity_data,
                policy_override=target_policy,
            )

        target_pk_value = getattr(target_obj, target_pk_col)

        if link_model is None:
            relation_items = cast(list[SQLModel], getattr(parent_obj, relation_name))
            if target_obj not in relation_items:
                relation_items.append(target_obj)
            continue

        stmt = select(link_model).where(
            getattr(link_model, link_parent_fk) == parent_pk_value,
            getattr(link_model, link_target_fk) == target_pk_value,
        )
        link_obj = (await ctx.session.exec(stmt)).first()
        _create_schema, link_update_schema = _flat_schemas(link_model)
        link_update_obj = _validate_model(link_update_schema, link_payload)
        if link_obj is None:
            if not link_policy.can_create(ctx.user, link_update_obj):
                raise ForbiddenError()
            link_data = {
                link_parent_fk: parent_pk_value,
                link_target_fk: target_pk_value,
                **link_update_obj.model_dump(exclude_unset=True),
            }
            ctx.session.add(link_model(**link_data))
            await ctx.session.flush()
        else:
            if not link_policy.can_update(ctx.user, link_obj, link_update_obj):
                raise ForbiddenError()
            for field, value in link_update_obj.model_dump(exclude_unset=True).items():
                setattr(link_obj, field, value)
            ctx.session.add(link_obj)
            await ctx.session.flush()


async def apply_nested_mutations(
    *,
    session: AsyncSession,
    user: User,
    model: type[SQLModel],
    db_obj: SQLModel,
    update_data: dict[str, object],
    nested_writes: NestedWriteConfig | None,
    nested_policy_registry: dict[type[SQLModel], CrudPolicy],
) -> None:
    """Apply nested mutations in-memory and flush without committing."""
    relation_policy_overrides = _relation_policy_overrides(
        model=model,
        nested_writes=nested_writes,
    )

    ctx = _MutationContext(
        session=session,
        user=user,
        nested_policy_registry=nested_policy_registry,
        relation_policy_overrides=relation_policy_overrides,
    )

    scalar_data, relation_data, relations = _split_payload(model, update_data)
    update_data.clear()
    update_data.update(scalar_data)
    for relation_name, relation_payload in relation_data.items():
        rel_info = relations[relation_name]
        await _apply_relation(
            ctx,
            parent_model=model,
            parent_obj=db_obj,
            relation_name=relation_name,
            rel_info=rel_info,
            relation_payload=relation_payload,
            parent_scalar_data=update_data,
        )


async def apply_nested_creates(
    *,
    session: AsyncSession,
    user: User,
    model: type[SQLModel],
    create_data: dict[str, object],
    nested_writes: NestedWriteConfig | None,
    nested_policy_registry: dict[type[SQLModel], CrudPolicy],
) -> SQLModel:
    """Create a full nested graph in-memory and flush without committing."""
    relation_policy_overrides = _relation_policy_overrides(
        model=model,
        nested_writes=nested_writes,
    )

    ctx = _MutationContext(
        session=session,
        user=user,
        nested_policy_registry=nested_policy_registry,
        relation_policy_overrides=relation_policy_overrides,
    )
    return await _create_model_graph(ctx, model, create_data)


def _relation_policy_overrides(
    *,
    model: type[SQLModel],
    nested_writes: NestedWriteConfig | None,
) -> dict[tuple[type[SQLModel], str], CrudPolicy]:
    relation_policy_overrides: dict[tuple[type[SQLModel], str], CrudPolicy] = {}
    if nested_writes:
        relations = inspect_relations(model)
        for relation_name, cfg in nested_writes.fields.items():
            rel_info = relations.get(relation_name)
            if rel_info is None:
                msg = f"Unknown relationship '{relation_name}' on model {model.__name__}"
                raise ConfigurationError(msg)
            if rel_info.target_model is not cfg.model:
                msg = (
                    f"Nested relation '{relation_name}' model mismatch: "
                    f"{rel_info.target_model.__name__} != {cfg.model.__name__}"
                )
                raise ConfigurationError(msg)
            if cfg.policy is not None:
                relation_policy_overrides[(model, relation_name)] = cast(
                    CrudPolicy, cfg.policy
                )
    return relation_policy_overrides
