from datetime import datetime, timedelta, timezone
from typing import Any, cast

from sqlalchemy import types
from sqlalchemy.engine.interfaces import Dialect
from sqlalchemy.sql.operators import OperatorType


class UTCDateTime(types.TypeDecorator[datetime]):
    """Store aware datetimes and return them in UTC."""

    impl = types.DateTime
    cache_ok = True

    def __init__(self) -> None:
        super().__init__(timezone=True)

    def __repr__(self) -> str:
        return "UTCDateTime()"

    def coerce_compared_value(
        self, op: OperatorType | None, value: Any
    ) -> types.TypeEngine[Any]:
        if isinstance(value, timedelta):
            return types.Interval()
        return self

    def process_bind_param(
        self, value: datetime | None, dialect: Dialect
    ) -> datetime | None:
        if value is None:
            return None
        if value.utcoffset() is None:
            raise ValueError(
                "Datetime values must have timezone information. "
                "Use datetime.now(timezone.utc), or annotate the field with "
                "NaiveDatetime for naive storage."
            )
        return value.astimezone(timezone.utc)

    def process_result_value(
        self, value: datetime | None, dialect: Dialect
    ) -> datetime | None:
        if value is None:
            return None
        if value.utcoffset() is None:
            # Databases without timezone support store UTC without an offset.
            return value.replace(tzinfo=timezone.utc)
        return value.astimezone(timezone.utc)


class AutoString(types.TypeDecorator):
    impl = types.String
    cache_ok = True
    mysql_default_length = 255

    def load_dialect_impl(self, dialect: Dialect) -> "types.TypeEngine[Any]":
        impl = cast(types.String, self.impl)
        if impl.length is None and dialect.name in {"mysql", "mariadb"}:
            return dialect.type_descriptor(types.String(self.mysql_default_length))
        return super().load_dialect_impl(dialect)