import json
import pickle
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Union

from .types import StrBytes


class Protocol(str, Enum):
    json = "json"
    pickle = "pickle"


def load_str_bytes(
    b: StrBytes,
    *,
    content_type: str = None,
    encoding: str = "utf8",
    proto: Protocol = None,
    allow_pickle: bool = False,
    json_loads: Callable[[str], Any] = json.loads,
) -> Any:
    if proto is None and content_type:
        if content_type.endswith(("json", "javascript")):
            pass
        elif allow_pickle and content_type.endswith("pickle"):
            proto = Protocol.pickle
        else:
            raise TypeError(f"Unknown content-type: {content_type}")

    proto = proto or Protocol.json

    if proto == Protocol.json:
        if isinstance(b, bytes):
            b = b.decode(encoding)
        return json_loads(b)
    elif proto == Protocol.pickle:
        if not allow_pickle:
            raise RuntimeError("Trying to decode with pickle with allow_pickle=False")
        bb = b if isinstance(b, bytes) else b.encode()
        return pickle.loads(bb)
    else:
        raise TypeError(f"Unknown protocol: {proto}")


def load_file(
    path: Union[str, Path],
    *,
    content_type: str = None,
    encoding: str = "utf8",
    proto: Protocol = None,
    allow_pickle: bool = False,
    json_loads: Callable[[str], Any] = json.loads,
) -> Any:
    path = Path(path)
    b = path.read_bytes()
    if content_type is None:
        if path.suffix in (".js", ".json"):
            proto = Protocol.json
        elif path.suffix == ".pkl":
            proto = Protocol.pickle

    return load_str_bytes(
        b,
        proto=proto,
        content_type=content_type,
        encoding=encoding,
        allow_pickle=allow_pickle,
        json_loads=json_loads,
    )
