import datetime
import json
from typing import Any

import sqlalchemy as sa
from pydantic import TypeAdapter
from sqlalchemy.engine import Dialect
from sqlalchemy_utils.types import ChoiceType as BaseChoiceType

from fansifter_common.adapters.encrypter import EncryptedStr
from fansifter_common.utils import timezone
from fansifter_common.utils.encoders import SafeJSONEncoder


class SafeJSONType(sa.types.TypeDecorator[str]):
    impl = sa.Text()
    hashable = False
    cache_ok = True

    def load_dialect_impl(self, dialect: Dialect) -> Any:
        return dialect.type_descriptor(self.impl_instance)

    def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
        if value is not None:
            value = json.dumps(value, cls=SafeJSONEncoder)
        return value

    def process_result_value(self, value: Any, dialect: Dialect) -> Any:
        if value is not None:
            value = json.loads(value)
        return value


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


class NaiveUTCDateTime(sa.types.TypeDecorator[datetime.datetime]):
    impl = sa.DateTime
    cache_ok = True

    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 = timezone.to_naive(value)
        return value

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


class PydanticType(sa.types.TypeDecorator):  # type: ignore[type-arg]
    impl = SafeJSONType
    cache_ok = True

    def __init__(self, type: Any, *, exclude_unset: bool = False) -> None:
        super().__init__()
        self.type = type
        self.type_adapter = TypeAdapter(type)
        self.exclude_unset = exclude_unset

    def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
        return self.type_adapter.dump_python(value, exclude_unset=self.exclude_unset)

    def process_result_value(self, value: Any, dialect: Dialect) -> Any:
        return self.type_adapter.validate_python(value)


class EncryptedStrType(sa.types.TypeDecorator[str]):
    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:  # noqa: ARG002
        return value

    def process_result_value(self, value: Any, dialect: Dialect) -> Any:  # noqa: ARG002
        if value is None:
            return value
        return EncryptedStr(value)

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