import base64
import pickle
from dataclasses import dataclass
from typing import Dict, Optional, Tuple

import safetensors.torch
import torch

from .aliases import PathOrStr

__all__ = [
    "state_dict_to_safetensors_file",
    "safetensors_file_to_state_dict",
]


@dataclass(eq=True, frozen=True)
class STKey:
    keys: Tuple
    value_is_pickled: bool


def encode_key(key: STKey) -> str:
    b = pickle.dumps((key.keys, key.value_is_pickled))
    b = base64.urlsafe_b64encode(b)
    return str(b, "ASCII")


def decode_key(key: str) -> STKey:
    b = base64.urlsafe_b64decode(key)
    keys, value_is_pickled = pickle.loads(b)
    return STKey(keys, value_is_pickled)


def flatten_dict(d: Dict) -> Dict[STKey, torch.Tensor]:
    result = {}
    for key, value in d.items():
        if isinstance(value, torch.Tensor):
            result[STKey((key,), False)] = value
        elif isinstance(value, dict):
            value = flatten_dict(value)
            for inner_key, inner_value in value.items():
                result[STKey((key,) + inner_key.keys, inner_key.value_is_pickled)] = inner_value
        else:
            pickled = bytearray(pickle.dumps(value))
            pickled_tensor = torch.frombuffer(pickled, dtype=torch.uint8)
            result[STKey((key,), True)] = pickled_tensor
    return result


def unflatten_dict(d: Dict[STKey, torch.Tensor]) -> Dict:
    result: Dict = {}

    for key, value in d.items():
        if key.value_is_pickled:
            value = pickle.loads(value.numpy().data)

        target_dict = result
        for k in key.keys[:-1]:
            new_target_dict = target_dict.get(k)
            if new_target_dict is None:
                new_target_dict = {}
                target_dict[k] = new_target_dict
            target_dict = new_target_dict
        target_dict[key.keys[-1]] = value

    return result


def state_dict_to_safetensors_file(state_dict: Dict, filename: PathOrStr):
    state_dict = flatten_dict(state_dict)
    state_dict = {encode_key(k): v for k, v in state_dict.items()}
    safetensors.torch.save_file(state_dict, filename)


def safetensors_file_to_state_dict(filename: PathOrStr, map_location: Optional[str] = None) -> Dict:
    if map_location is None:
        map_location = "cpu"
    state_dict = safetensors.torch.load_file(filename, device=map_location)
    state_dict = {decode_key(k): v for k, v in state_dict.items()}
    return unflatten_dict(state_dict)