from abc import ABC, abstractmethod
from typing import Any, Type, Union, Annotated, get_origin, get_args, List, Optional
from uuid import UUID as _UUID

from bson import Binary
from bson import ObjectId as _ObjectId
from pydantic import GetCoreSchemaHandler, BaseModel
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema
from pydantic.functional_serializers import PlainSerializer


def is_optional(field):
    return get_origin(field) is Union and \
           type(None) in get_args(field)


class ObjectIdPydanticAnnotation:
    @classmethod
    def validate_object_id(cls, v: Any, handler) -> _ObjectId:
        if isinstance(v, _ObjectId):
            return v

        s = handler(v)
        if _ObjectId.is_valid(s):
            return _ObjectId(s)
        else:
            raise ValueError("Invalid ObjectId")

    @classmethod
    def __get_pydantic_core_schema__(cls, source_type, _handler) -> core_schema.CoreSchema:
        assert source_type is _ObjectId or (is_optional(source_type) and _ObjectId in get_args(source_type))
        return core_schema.no_info_wrap_validator_function(
            cls.validate_object_id,
            core_schema.str_schema(),
            serialization=core_schema.to_string_ser_schema(),
        )

    @classmethod
    def __get_pydantic_json_schema__(cls, _core_schema, handler) -> JsonSchemaValue:
        return handler(core_schema.str_schema())


ObjectId = Annotated[_ObjectId, ObjectIdPydanticAnnotation]


class Uuid(_UUID):
    @classmethod
    def validate_uuid(cls, v: Any) -> _UUID:
        if isinstance(v, _UUID): return v
        if isinstance(v, str): return _UUID(v)
        if isinstance(v, Binary): return v.as_uuid()

    # def to_binary(self):
    #     return Binary.from_uuid(self)

    @classmethod
    def from_uuid(cls, payload) -> Binary:
        if isinstance(payload, Binary): return payload
        return Binary.from_uuid(payload)

    @classmethod
    def __get_pydantic_core_schema__(
            cls, source: Type[Any], handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        return core_schema.no_info_after_validator_function(
            cls.validate_uuid,
            core_schema.uuid_schema(),
            serialization=core_schema.plain_serializer_function_ser_schema(
                function=cls.from_uuid
            ),
        )

# add json plain serializer to Uuid
# Uuid = Annotated[
#     Uuid, PlainSerializer(lambda x: str(x), return_type=str, when_used='json')
# ]

MaybeUuid = Union[Uuid, _UUID, Binary]


def to_binary_uuid(uuid_: Union[MaybeUuid, List[MaybeUuid]]):
    if isinstance(uuid_, list): return [to_binary_uuid(u) for u in uuid_]
    if isinstance(uuid_, Binary): return uuid_
    return Binary.from_uuid(uuid_)


def serialize_object_ids_in_dict(mey_be_dict: Any):
    if not isinstance(mey_be_dict, dict): return
    for key, val in mey_be_dict.items():
        if isinstance(val, _ObjectId):
            mey_be_dict[key] = str(val)
        elif isinstance(val, dict):
            serialize_object_ids_in_dict(val)
        elif isinstance(val, list):
            for item in val:
                serialize_object_ids_in_dict(item)


class DbSettingsABC(ABC):

    @property
    @abstractmethod
    def qdrant_host(self) -> Optional[str]: ...

    @property
    @abstractmethod
    def qdrant_port(self) -> Optional[int]: ...

    @property
    @abstractmethod
    def mongo_db_uri(self) -> Optional[str]: ...

    @property
    @abstractmethod
    def redis_uri(self) -> Optional[str]: ...

    @property
    @abstractmethod
    def sql_write_uri(self) -> Optional[str]: ...

    @property
    @abstractmethod
    def sql_read_uri(self) -> Optional[str]: ...

    @property
    @abstractmethod
    def sql_uri(self) -> Optional[str]: ...


class QueryBase(BaseModel):
    offset: Optional[int] = None
    limit: Optional[int] = None
    filters: Optional[Union[dict, List[Any]]] = None
    order_by: Optional[Any] = None

    id: Optional[Union[ObjectId, Uuid, str]] = None
    ids: Optional[List[Union[ObjectId, Uuid, str]]] = None

    @classmethod
    def loop(cls, func, *args, **kwargs) -> List[Any]:
        page = 0
        page_size = 100
        res = []
        has_more = True
        base_query: QueryBase = kwargs.get("query")
        if 'query' in kwargs: del kwargs['query']
        base_query_dict: dict = base_query.model_dump() if base_query is not None else {}
        if 'offset' in base_query_dict: del base_query_dict['offset']
        if 'limit' in base_query_dict: del base_query_dict['limit']
        while page < 10000 and has_more:
            query = QueryBase(**base_query_dict, limit=page_size, offset=page * page_size)

            output = func(*args, **kwargs, query=query)
            if len(output) == 0: break
            has_more = len(output) == page_size
            res += output
            page += 1
        return res
