from typing import Any, Self, Required, TypedDict, Literal, Dict
from uuid import UUID as _DEFAULT_UUID

try:
    from uuid_utils import UUID as _UUID
    from uuid_utils import uuid7, uuid4
except ModuleNotFoundError as e:  # pragma: no cover
    raise RuntimeError(
        'The `ulid` module requires "uuid_utils" to be installed. '
        'You can install it with "pip install python-ulid".'
    ) from e
from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import PydanticCustomError, SchemaSerializer, core_schema

# from pydantic import UUID1
from loguru import logger


class UuidSchema(TypedDict, total=False):
    type: Required[Literal["uuid"]]
    version: Literal[1, 3, 4, 5, 6, 7]
    strict: bool
    ref: str
    metadata: Dict[str, Any]
    serialization: core_schema.SerSchema


class UUID(_UUID):
    def encode(self) -> str:
        return str(self)

    # def __get_pydantic_json_schema__(
    #     self,
    #     core_schema: core_schema.CoreSchema,
    #     handler: GetJsonSchemaHandler | None = None,
    # ) -> JsonSchemaValue:

    #     field_schema = handler(core_schema)
    #     field_schema.pop("anyOf", None)  # remove the bytes/str union
    #     field_schema.update(type="string", format=f"uuid{self.uuid_version}")
    #     return field_schema

    @classmethod
    def __get_pydantic_core_schema__(
        cls,
        source: type[Any],
        handler: GetCoreSchemaHandler,
    ) -> core_schema.CoreSchema:
        schema = core_schema.no_info_after_validator_function(
            cls._validate_ulid,
            core_schema.union_schema(
                [
                    core_schema.is_instance_schema(_DEFAULT_UUID),
                    core_schema.is_instance_schema(_UUID),
                    core_schema.is_instance_schema(cls),
                    core_schema.int_schema(),
                    core_schema.bytes_schema(),
                    core_schema.str_schema(),
                ]
            ),
            serialization=core_schema.plain_serializer_function_ser_schema(
                str,
                info_arg=False,
                return_schema=core_schema.str_schema(),
                when_used="json",
            ),
        )
        cls.__pydantic_serializer__ = SchemaSerializer(
            schema
        )  # <-- this is necessary for pydantic-core to serialize

        return schema

    @classmethod
    def _validate_ulid(
        cls,
        value: Any,
        handler: core_schema.ValidatorFunctionWrapHandler | None = None,
        info: core_schema.ValidationInfo | None = None,
    ) -> Any:
        try:
            if isinstance(value, int):
                ulid = cls(int=value)
            elif isinstance(value, str):
                ulid = cls(hex=value)
            elif isinstance(value, cls) or isinstance(value, _DEFAULT_UUID):
                ulid = cls(int=value.int)
            elif hasattr(value, "int"):
                ulid = cls(int=value.int)
            elif isinstance(value, bytes):
                ulid = cls(bytes=value)
            else:
                raise TypeError(f"Invalid type: {type(value)}")
        except ValueError as e:
            raise PydanticCustomError(
                "uuid_format",
                "Unrecognized format",
            ) from e
        return handler(ulid) if handler else str(ulid)

    @classmethod
    def create_v7(cls) -> Self:
        return cls(int=uuid7().int)
