from datetime import UTC, datetime
from typing import Any

from sqlalchemy import DateTime
from sqlalchemy.engine import Dialect
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.types import TypeDecorator


class UTCDateTime(TypeDecorator[datetime]):
    """`DateTime(timezone=True)` does NOT round-trip tzinfo on SQLite — confirmed empirically:
    a tz-aware datetime goes in, a naive one comes back out, and every
    `expires_at < datetime.now(UTC)` style comparison in auth/service.py then raises
    `TypeError: can't compare offset-naive and offset-aware datetimes`. Not a quirk worth a
    per-column workaround — every timestamp in this app is UTC by convention (docs/PLAN.md), so
    re-attach UTC on load here, once, rather than trust the driver to preserve it.
    """

    impl = DateTime(timezone=True)
    cache_ok = True

    def process_bind_param(self, value: datetime | None, dialect: Dialect) -> datetime | None:
        if value is not None and value.tzinfo is None:
            raise ValueError(
                "naive datetime passed to a UTCDateTime column — always construct with "
                "datetime.now(UTC), never datetime.now() or datetime.utcnow()"
            )
        return value

    def process_result_value(self, value: Any, dialect: Dialect) -> datetime | None:
        if value is not None and value.tzinfo is None:
            value = value.replace(tzinfo=UTC)
        return value  # type: ignore[no-any-return]


class Base(DeclarativeBase):
    # Every timestamp in this schema is UTC (docs/PLAN.md's convention) — without this, a bare
    # `Mapped[datetime]` infers a plain DateTime, which hits the SQLite round-trip bug above on
    # every single column instead of being fixed once, here, for every current and future model.
    type_annotation_map = {
        datetime: UTCDateTime(),
    }