from __future__ import annotations

from uuid import UUID, uuid4

from tigrbl.runtime.status.exceptions import HTTPException

from tigrbl.decorators.hook import hook_ctx
from tigrbl.specs import IO, F, S, Pair, ForeignKeySpec, acol
from tigrbl.orm.mixins import BulkCapable, Replaceable
from tigrbl.orm.tables import Base
from tigrbl.types import (
    Integer,
    LargeBinary,
    PgUUID,
    SAEnum,
    UniqueConstraint,
    relationship,
    Mapped,
)
from swarmauri_core.crypto.types import ExportPolicy, KeyUse
from swarmauri_core.key_providers.types import (
    KeySpec,
    KeyClass,
    KeyAlg as ProviderKeyAlg,
)

from ..utils import b64d
from .key import Key, KeyAlg


class KeyVersion(Base, BulkCapable, Replaceable):
    __tablename__ = "key_versions"
    __resource__ = "key_version"
    __table_args__ = (UniqueConstraint("key_id", "version"),)

    # Provide explicit primary key to avoid mixin materialization issues in tests
    id: Mapped[UUID] = acol(
        storage=S(
            type_=PgUUID(as_uuid=True),
            primary_key=True,
            default=uuid4,
        ),
        field=F(py_type=UUID),
        io=IO(out_verbs=("read", "list")),
    )

    key_id: Mapped[UUID] = acol(
        storage=S(
            type_=PgUUID(as_uuid=True),
            fk=ForeignKeySpec(target="keys.id", on_delete="CASCADE"),
            nullable=False,
            index=True,
        ),
        field=F(py_type=UUID, required_in=("create",)),
        io=IO(
            in_verbs=("create",),
            out_verbs=("read", "list"),
            filter_ops=("eq",),
        ),
    )

    version: Mapped[int] = acol(
        storage=S(type_=Integer, nullable=False),
        field=F(py_type=int, required_in=("create",)),
        io=IO(in_verbs=("create",), out_verbs=("read", "list"), sortable=True),
    )

    status: Mapped[str] = acol(
        storage=S(
            type_=SAEnum(
                "active",
                name="VersionStatus",
                native_enum=True,
                validate_strings=True,
                create_constraint=True,
            ),
            nullable=False,
        ),
        field=F(py_type=str, required_in=("create",)),
        io=IO(in_verbs=("create",), out_verbs=("read", "list")),
    )

    def _pair_material(ctx):
        payload = ctx.get("payload") or {}
        raw = payload.get("public_material_b64")
        if raw is None:
            return Pair(raw=None, stored=None)
        return Pair(raw=raw, stored=b64d(raw))

    public_material: Mapped[bytes | memoryview | None] = acol(
        storage=S(type_=LargeBinary, nullable=True),
        io=IO().paired(
            _pair_material, alias="public_material_b64", verbs=("create", "update")
        ),
    )

    key = relationship("Key", back_populates="versions", lazy="joined")

    @hook_ctx(ops=("create", "bulk_create"), phase="PRE_HANDLER")
    async def _coerce_key_id_uuid(cls, ctx):
        payload = ctx.get("payload")
        if payload is None:
            return

        def _coerce_item(item):
            if not isinstance(item, dict):
                return
            key_id = item.get("key_id")
            if key_id is None or isinstance(key_id, UUID):
                return
            try:
                item["key_id"] = UUID(str(key_id))
            except Exception as exc:
                raise HTTPException(status_code=400, detail="Invalid key_id") from exc

        if isinstance(payload, list):
            for item in payload:
                _coerce_item(item)
            return

        if isinstance(payload, dict):
            _coerce_item(payload)
            items = payload.get("items")
            if isinstance(items, list):
                for item in items:
                    _coerce_item(item)

    @hook_ctx(ops="create", phase="POST_HANDLER")
    async def _generate_material(cls, ctx):
        obj = ctx.get("result")
        if obj is None or obj.public_material is not None:
            return
        key_provider = ctx.get("key_provider")
        if key_provider is None:
            raise HTTPException(status_code=500, detail="Key provider missing")
        db = ctx.get("db")
        if db is None:
            raise HTTPException(status_code=500, detail="DB session missing")
        key_id = getattr(obj, "key_id", None)
        if key_id is None:
            raise HTTPException(status_code=400, detail="Missing key_id")
        key_obj = await db.get(Key, UUID(str(key_id)))
        if key_obj is None:
            raise HTTPException(status_code=404, detail="Key not found")
        if key_obj.algorithm in (KeyAlg.AES256_GCM, KeyAlg.CHACHA20_POLY1305):
            material = await key_provider.random_bytes(32)
        else:  # pragma: no cover - defensive
            raise HTTPException(status_code=400, detail="Unsupported algorithm")
        await key_provider.import_key(
            KeySpec(
                klass=KeyClass.symmetric,
                alg=ProviderKeyAlg.AES256_GCM,
                uses=(KeyUse.ENCRYPT, KeyUse.DECRYPT),
                export_policy=ExportPolicy.SECRET_WHEN_ALLOWED,
                label=str(key_id),
            ),
            material,
        )
        obj.public_material = material

    @hook_ctx(ops="create", phase="POST_RESPONSE")
    async def _scrub_material(cls, ctx):
        """Remove raw key material from the response payload."""
        obj = ctx.get("result")
        if obj is None:
            return
        if isinstance(obj, dict):
            obj.pop("public_material", None)
        else:
            obj.public_material = None
