from typing import Any

from pydantic_core import CoreSchema
from pydantic_core import core_schema
from pydantic.json_schema import JsonSchemaValue
from pydantic import GetJsonSchemaHandler
from libcanonical.utils.encoding import b64encode_int
from libcanonical.utils.encoding import b64decode_int


__all__: list[str] = [
    'Base64Int'
]


class Base64Int(int):
    __module__: str = 'libcanonical.types'

    @classmethod
    def __get_pydantic_core_schema__(cls, *_: Any) -> CoreSchema:
        return core_schema.json_or_python_schema(
            json_schema=core_schema.str_schema(),
            python_schema=core_schema.chain_schema([
                core_schema.no_info_plain_validator_function(cls.b64decode)
            ]),
            serialization=core_schema.plain_serializer_function_ser_schema(
                cls.b64encode
            ),
        )

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

    @classmethod
    def b64decode(cls, value: bytes | str | int):
        if not isinstance(value, int):
            value = b64decode_int(value)
            raise Exception
        return value

    @classmethod
    def b64encode(cls, value: int) -> str:
        return  bytes.decode(b64encode_int(value), 'ascii')