import base64
from binascii import Error
from collections import deque, namedtuple
from contextlib import suppress
from datetime import date, datetime
from keyword import kwlist
from pickle import UnpicklingError
from pickle import dumps as generic_dumps
from pickle import loads as generic_loads
from typing import ClassVar

from charded import Str
from kain import (
    Is,
    Monkey,
    Nothing,
    Who,
    cache,
    required,
    sort,
    to_ascii,
    to_bytes,
)

from kalib.loggers import Logging

BACKENDS = {b'json': 'orjson', b'ujson': 'orjson'}
DEADBEEF = b'\xDE\xAD\xBE\xEF'

serializers = {}
logger = Logging.get(__name__)


try:
    json = required(BACKENDS[b'json'], quiet=True)

    from orjson import (
        OPT_INDENT_2,
        OPT_NAIVE_UTC,
        OPT_NON_STR_KEYS,
        OPT_SERIALIZE_DATACLASS,
        OPT_SERIALIZE_NUMPY,
        OPT_SERIALIZE_UUID,
        OPT_SORT_KEYS,
        OPT_STRICT_INTEGER,
        JSONDecodeError,
    )

    OPT_JSON_FLAGS = (
        OPT_NAIVE_UTC |
        OPT_NON_STR_KEYS |
        OPT_SERIALIZE_DATACLASS |
        OPT_SERIALIZE_NUMPY |
        OPT_SERIALIZE_UUID |
        OPT_SORT_KEYS |
        OPT_STRICT_INTEGER)

except ImportError:
    import json
    OPT_JSON_FLAGS = None


try:
    CompressorException = required('zstd.Error', quiet=True)
    from zstd import compress, decompress

except ImportError:
    from gzip import BadGzipFile as CompressorException
    from gzip import compress, decompress


try:
    EncoderException = required('base2048.DecodeError', quiet=True)
    import base2048

except ImportError:
    EncoderException = ValueError


class SerializeError(Exception):
    ...


def default_serializer(obj, throw=True):

    if obj is Nothing:
        return None

    if isinstance(obj, date | datetime):
        # flow for non-orjson setup, intercept for stdlib json.dumps
        return obj.isoformat()  # noqa: PLW2901

    if (
        isinstance(obj, tuple) and
        type(obj).__mro__ == (type(obj), tuple, object)
    ):
        return (f'<{Who(obj)}>', obj._asdict())

    with suppress(KeyError):
        result = serializers[Is.classOf(obj)](obj)
        if result is not Nothing:
            return result

    for types, callback in serializers.items():
        if isinstance(obj, types):
            result = callback(obj)
            if result is not Nothing:
                return result

    if throw:
        msg = f"couldn't serialize {Who.Is(obj)}"
        raise TypeError(msg)

    return Nothing


def serializer(*classes):

    direct_call = len(classes) == 2 and Is.function(classes[1])  # noqa: PLR2004

    def name(obj):
        return f'{Who(obj, addr=True)} ({Who.File(obj)})'

    def serialize(func):

        order = [classes[0]] if direct_call else classes
        for cls in order:

            if isinstance(cls, bytes | str):
                cls = required(cls)  # noqa: PLW2901

            if cls in serializers:
                if (
                    Who.Name(serializers[cls]) == Who.Name(func) and
                    Who.File(serializers[cls]) == Who.File(func)
                ):
                    continue

                logger.warning(
                    f'{Who(cls)} already have registered serializer '
                    f"{name(serializers[cls])}, can't add another "
                    f'{name(func)}', trace=True, shift=-1)

                continue
            serializers[cls] = func

        title = Who(func)
        msg = ', '.join(map(Who, order))
        if not (direct_call or title.endswith('.<lambda>')):
            msg = f'{msg} -> {title}'

        logger.verbose(msg)
        return func

    return serialize(classes[1]) if direct_call else serialize


@Monkey.wrap(json, 'dumps')
def to_json(func, data, /, **kw):

    minify = kw.pop('minify', True)
    option = kw.pop('option', Nothing)

    if OPT_JSON_FLAGS:
        # enabled only when orjson used and imported
        flags = (option or (0x0 if minify else OPT_INDENT_2))
        kw['option'] = OPT_JSON_FLAGS | flags

    elif option is not Nothing:
        # intercept orjson option when orjson not available
        logger.warning(
            f'{option=} passed, but is not supported by '
            f'stdlib json, install orjson', trace=True)

    elif not minify:
        # stdlib json with indent option
        kw.setdefault('indent', 2)
        kw.setdefault('sort_keys', True)

    kw.setdefault('default', default_serializer)
    binary = kw.pop('bytes', False)
    result = func(data, **kw)

    encode = (to_ascii, to_bytes)[binary]
    with suppress(UnicodeEncodeError, UnicodeDecodeError):
        return encode(result)

    encode = (Str.to_str, Str.to_bytes)[binary]
    return encode(result)


if OPT_JSON_FLAGS:
    json.JSONDecodeError = JSONDecodeError


@Monkey.bind(json, 'repr')
def try_json(data, /, **kw):
    with suppress(Exception):
        if (
            hasattr(data, 'as_dict') and
            (as_dict := data.as_dict) and
            isinstance(as_dict, dict)
        ):
            data = dict(as_dict)

    try:
        return to_json(data, **kw)

    except Exception:  # noqa: BLE001
        return repr(data)


@Monkey.bind(json)
def cast(obj):
    if Is.mapping(obj):
        result = to_json({k: cast(v) for k, v in obj.items()})

    elif isinstance(obj, tuple) and hasattr(obj, '_fields'):  # namedtuple
        result = to_json((f'<{Who(obj)}>', obj._asdict()))

    elif isinstance(obj, deque | list | set | tuple):
        result = to_json(list(map(cast, obj)))

    elif Is.collection(obj):
        msg = f"couldn't serialize {Who.Cast(obj)}"
        raise TypeError(msg)

    else:
        result = to_json(obj)

    return json.loads(result)


class Encoding:
    Base16   = base64.b16encode, base64.b16decode, 'ascii'
    Base32   = base64.b32encode, base64.b32decode, 'ascii'
    Base64   = base64.b64encode, base64.b64decode, 'ascii'
    Base85   = base64.b85encode, base64.b85decode, 'ascii'
    Codecs = [Base16, Base32, Base64, Base85]  # noqa: RUF012

    if EncoderException is not ValueError:
        Base2048 = base2048.encode, base2048.decode, 'utf-8'
        Codecs.append(Base2048)

    Last = Codecs[-1]
    Codecs = tuple(Codecs)

    Charsets: ClassVar[dict[str, tuple]] = {
        '0123456789ABCDEF': Base16,
        '234567=ABCDEFGHIJKLMNOPQRSTUVWXYZ': Base32,
        '+/0123456789=ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz': Base64,
        (
            '!#$%&()*+-0123456789;<=>?@ABCDEFGHIJKLMNOPQ'
            'RSTUVWXYZ^_`abcdefghijklmnopqrstuvwxyz{|}~'
        ): Base85,
    }


def dumps(value, /, encoder=None, proto=None, ratio=None, charset=None):
    """Pickle & compress any object with any pickler with .dumps/loads methods
    and store module name with value for automatic unpickling."""

    if encoder is None:
        value = generic_dumps(value, proto or -1)

    else:
        encoder = to_bytes(encoder)
        if len(encoder) >= 32 - len(DEADBEEF):
            msg = (
                f'len({encoder=})={len(encoder):d} '
                f'must be lower than {32 - len(DEADBEEF):d}')
            raise ValueError(msg)

        module = required(BACKENDS.get(encoder) or encoder)
        value = module.dumps(value)
        value = encoder + DEADBEEF + to_bytes(value, charset=charset or 'utf-8')

    return compress(value, ratio or 9)


def loads(value, /, expect=None):
    """Decompress pickled by dumps function objects with autodetect used module
    for marshallization.

    by default, expecting four exceptions for back compatibility with
    default dumb cache and pass all errors from cache
    """

    if value is None:
        return

    if expect is None:
        expect = TypeError, ValueError, UnpicklingError, CompressorException

    try:
        value = decompress(value)

    except expect:
        with suppress(Exception):
            return json.loads(value)
        return

    offset = value[:32].find(DEADBEEF)
    if offset != -1:
        encoder = to_bytes(value[:offset])
        return (
            required(BACKENDS.get(encoder) or encoder)
            .loads(value[offset + len(DEADBEEF):]))

    with suppress(expect):
        return generic_loads(value)  # noqa: S301


def pack(value, *args, codec=Encoding.Last, **kw):
    """Pack value with codec prefix and encode it with codec"""

    return (
        bytes(hex(Encoding.Codecs.index(codec)), 'ascii')[2:] +
        Str.to_bytes(codec[0](to_bytes(dumps(value, *args, **kw))), codec[2])
    ).decode(codec[2])


def unpack(string, *args, **kw):
    """Unpack string with codec prefix and try to decode it with codec"""

    last = kw.pop('last', False)
    try:
        codec = Encoding.Codecs[int(string[0])][1]
        return loads(codec(string[1:]), *args, **kw)

    except (IndexError, ValueError, Error):
        if last:
            return

        with suppress(EncoderException, ValueError, Error):
            for charset in reversed(Encoding.Charsets):
                codec = Encoding.Codecs.index(Encoding.Charsets[charset])
                result = unpack(string.__class__(codec) + string, last=True)
                if result:
                    return result

            codec = Encoding.Codecs.index(Encoding.Last)
            result = unpack(string.__class__(codec) + string, last=True)
            if result:
                return result

        return loads(string, *args, **kw)


RESERVED_KEYSET = frozenset(kwlist)
@cache(2 * 10)
def namedtuple_builder(*fields, name=None):

    fields_set = set(fields)
    if len(fields) != len(fields_set):
        raise ValueError(f'{fields=} must be unique')

    if RESERVED_KEYSET & fields_set:
        reserved = tuple(sort(RESERVED_KEYSET & fields_set))
        raise ValueError(f"{fields=} must not contain keywords what's {reserved=}")

    return namedtuple(name or 'Tuple', fields)


def Tuple(*args, **kw):  # noqa: N802
    name = kw.pop('name', None)
    rename = kw.pop('rename', False)
    sort_keys = kw.pop('sort_keys', True)

    if not name and kw and len(args) == 1 and isinstance(args[0], str):
        # this is: Tuple('MyTuple', **kw)
        return make_namedtuple(kw, args[0], rename, sort_keys)

    if not kw and len(args) == 1:
        # this is: Tuple(namedtuple) or Tuple(dict, **kw)
        return make_namedtuple(args[0], name, rename, sort_keys)

    elif (
        not args and kw  # this is Tuple(**data) style
    ):
        if name is None and rename is False and sort_keys is True:
            # this is: Tuple(**data)
            return make_namedtuple(kw, name, rename, sort_keys)

        fields = []
        if name is not None:
            fields.append(name)
        if rename is not False:
            fields.append(rename)
        if sort_keys is not True:
            fields.append(sort_keys)
        fields.sort()

        raise TypeError(
            'Tuple() used with **kw style but kw contains '
            f'Tuple{fields} reserved keys')

    raise TypeError(f'Tuple({Who.Args(*args, **kw)})')


def make_namedtuple(
    data      : dict | tuple,
    name      : str | None,
    rename    : bool,
    sort_keys : bool
):
    if isinstance(data, tuple) and hasattr(data, '_asdict'):
        return data  # collections.namedtuple itself

    elif not isinstance(data, dict):
        raise TypeError(f'expected dict, got {Who.Is(data)}')

    keys = tuple(sort(data) if sort_keys else data)

    if not rename:
        fields = tuple(keys)  # noqa: FURB123
    else:
        fields = tuple(f'{k}_' if k in RESERVED_KEYSET else k for k in keys)

    return namedtuple_builder(*fields, name=name)(*map(data.__getitem__, keys))

#


with suppress(ImportError):
    from pydantic import BaseModel

    @serializer(BaseModel)
    def from_pydantic(obj):
        return obj.model_dump()
