from __future__ import annotations

import base64
import json
import pickle  # nosec: B301,B403
from typing import Any

from tortoise.indexes import Index


class JsonEncoder(json.JSONEncoder):
    def default(self, obj: Any) -> Any:
        if isinstance(obj, Index):
            if hasattr(obj, "describe"):
                # For tortoise>=0.24
                return obj.describe()
            return {
                "type": "index",
                "val": base64.b64encode(pickle.dumps(obj)).decode(),  # nosec: B301
            }
        else:
            return super().default(obj)


def object_hook(obj: dict[str, Any]) -> Any:
    if (type_ := obj.get("type")) and type_ == "index" and (val := obj.get("val")):
        return pickle.loads(base64.b64decode(val))  # nosec: B301
    return obj


def load_index(obj: dict[str, Any]) -> Index:
    """Convert a dict that generated by `Index.decribe()` to a Index instance"""
    try:
        index = Index(fields=obj["fields"] or obj["expressions"], name=obj.get("name"))
    except KeyError:
        return object_hook(obj)
    if extra := obj.get("extra"):
        index.extra = extra
    if idx_type := obj.get("type"):
        index.INDEX_TYPE = idx_type
    return index


def encoder(obj: dict[str, Any]) -> str:
    return json.dumps(obj, cls=JsonEncoder)


def decoder(obj: str | bytes) -> Any:
    return json.loads(obj, object_hook=object_hook)
