import datetime
from typing import Any

import sqlalchemy as sa
from audience_common import datetimeutc
from sqlalchemy.engine import Dialect
from sqlalchemy_utils.types import ChoiceType as BaseChoiceType

from campaigns.core.types import EncryptedToken


class NaiveUTCDateTime(sa.TypeDecorator):  # type: ignore[type-arg]
    impl = sa.DateTime()
    cache_ok = True

    @property
    def python_type(self) -> Any:
        return self.impl.python_type

    def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
        if isinstance(value, datetime.datetime):
            if not value.tzinfo:
                raise TypeError("tzinfo is required")
            value = datetimeutc.to_naive(value)
        return value

    def process_result_value(self, value: Any, dialect: Dialect) -> Any:
        if isinstance(value, datetime.datetime):
            value = datetimeutc.to_aware(value)
        return value

    def process_literal_param(self, value: Any, dialect: Dialect) -> Any:
        return value


class EncryptedTextType(sa.TypeDecorator):  # type: ignore[type-arg]
    impl = sa.Text()
    cache_ok = True

    @property
    def python_type(self) -> Any:
        return self.impl.python_type

    def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
        return value

    def process_result_value(self, value: Any, dialect: Dialect) -> Any:
        if value is None:
            return value
        return EncryptedToken(value)

    def process_literal_param(self, value: Any, dialect: Dialect) -> Any:
        return value


class ChoiceType(BaseChoiceType):  # type: ignore[misc]
    cache_ok = True
    impl = sa.Text()

    def process_literal_param(self, value: Any, dialect: Dialect) -> Any:
        return value