from __future__ import annotations

import collections
import os
import sys
import tempfile
from typing import Any, Union

import numpy as np
import torch
from loguru._logger import Core as _Core
from loguru._logger import Logger as _Logger

import jammy.io as jio
from jammy.logging.logger import LOG_FORMAT
from jammy.utils.registry import CallbackRegistry
from jamtorch.logging.logger import _rank0_only_filter


def _default_replayer_fallback(obj, *args, **kwargs):
    del args, kwargs
    if obj is type(None):
        return None
    item_type = type(obj)
    raise ValueError(f'Unknown item type: "{item_type}".')


class _RePlayerCallbackRegistry(CallbackRegistry):
    def dispatch(self, name, *args, **kwargs):
        return super().dispatch_direct(name, *args, **kwargs)


def replay_proc_dict(obj: dict[str, Any]):
    rtn_dict = {}
    for k, v in obj.items():
        rtn_dict[k] = replayer_processing(v)
    return rtn_dict


def replay_proc_module(obj: torch.nn.Module):
    grads_dict = {}
    for k, v in obj.named_parameters():
        # record weights value and grad if exists
        if v.grad is not None:
            grads_dict[k] = v.grad.clone().detach()
    return {
        "weights": obj.state_dict(),
        "grads": grads_dict,
    }


def replay_proc_basic(obj: list[Union[bool, int, float, str, complex]]):
    assert type(obj) in [bool, int, float, str, complex, np.ndarray, torch.Tensor]
    if isinstance(obj, torch.Tensor):
        return obj.cpu()
    return obj


def replay_proc_iter(obj: Union[list[Any], tuple[Any, ...]]):
    rtn = []
    for cur_obj in obj:
        rtn.append(replayer_processing(cur_obj))
    return rtn


_replayer_registry = _RePlayerCallbackRegistry()
_replayer_registry.set_fallback_callback(_default_replayer_fallback)
_replayer_registry.register(bool, replay_proc_basic)
_replayer_registry.register(int, replay_proc_basic)
_replayer_registry.register(str, replay_proc_basic)
_replayer_registry.register(float, replay_proc_basic)
_replayer_registry.register(complex, replay_proc_basic)
_replayer_registry.register(np.ndarray, replay_proc_basic)
_replayer_registry.register(torch.Tensor, replay_proc_basic)
_replayer_registry.register(dict, replay_proc_dict)
_replayer_registry.register(collections.OrderedDict, replay_proc_dict)
_replayer_registry.register(collections.defaultdict, replay_proc_dict)
_replayer_registry.register(tuple, replay_proc_iter)
_replayer_registry.register(list, replay_proc_iter)


def replayer_processing(obj: Any):
    return _replayer_registry.dispatch(type(obj), obj)


class RePlayer:
    _KEY_: str = "jam_replay"

    def __init__(self, savedir, level="INFO"):
        self.savedir = savedir
        self.dump_counter = collections.defaultdict(int)
        self.data = {}
        if not os.path.exists(savedir):
            os.makedirs(savedir)
        self.logger = _Logger(
            core=_Core(),
            exception=None,
            depth=0,
            record=False,
            lazy=False,
            colors=False,
            raw=False,
            capture=True,
            patchers=[],
            extra={},
        )
        self.logger.add(
            os.path.join(savedir, f"_{RePlayer._KEY_}.log"),
            level=level,
            mode="w",
            format=LOG_FORMAT,
            filter=_rank0_only_filter,
        )
        self.logger.add(
            sys.stdout,
            level="INFO",
            format=LOG_FORMAT,
            filter=_rank0_only_filter,
        )

    def info(self, message, rank0_only=False):
        self.logger.opt(depth=1).bind(rank0_only=rank0_only).info(message)

    def debug(self, message, rank0_only=False):
        self.logger.opt(depth=1).bind(rank0_only=rank0_only).debug(message)

    def critical(self, message, rank0_only=False):
        self.logger.opt(depth=1).bind(rank0_only=rank0_only).critical(message)

    def log(self, **kwargs):
        self.data.update(replayer_processing(kwargs))

    def reset(self):
        self.data = {}
        self.dump_counter = collections.defaultdict(int)

    def dump(self):
        self.dump_counter[RePlayer._KEY_] += 1
        save_fp = f"{self.savedir}/{RePlayer._KEY_}_{self.dump_counter[RePlayer._KEY_]:04d}.pkl"
        jio.dump(
            save_fp,
            (self.data, self.dump_counter),
        )
        self.data = {}
        self.info(f"Dumped {RePlayer._KEY_} at {save_fp}")

    def dump_state(self, key: str, value: Any, is_count: bool = False):
        if is_count:
            self.dump_counter[key] += 1
            jio.dump(
                f"{self.savedir}/{key}_{self.dump_counter[key]:04d}.pkl",
                value,
            )
        else:
            jio.dump(
                f"{self.savedir}/{key}.pkl",
                value,
            )

    def load(self, idx, savedir=None):
        if savedir is None:
            savedir = self.savedir
        self.data, self.dump_counter = jio.load(
            f"{savedir}/jam_replay_{idx:04d}.pkl",
        )

    def index(self):
        return self.dump_counter[RePlayer._KEY_]


_default_replayer = None


def get_default_replayer(path: str | None = None):
    global _default_replayer
    if _default_replayer is None:
        if path is None:
            path = tempfile.mkdtemp(prefix="jam_replay_")
        _default_replayer = RePlayer(savedir=path)
    return _default_replayer
