"""
hegm.backends.pytorch -- PyTorch backend.

Handles monkey-patching of ``torch.nn.Module`` and ``torch.optim.Optimizer``,
checkpoint save/load, RNG state management, and GPU resource cleanup.
"""

import functools
import os
import sys
import threading
import weakref
from typing import Any, Dict, List, Set

from hegm._config import CHECKPOINT_PATH, log, warn
from hegm.backends import AbstractBackend, register

# ---------------------------------------------------------------------------
# Tracking state
# ---------------------------------------------------------------------------

_tracked_models: List[weakref.ref] = []
_tracked_optimizers: List[weakref.ref] = []

_checkpoint_loaded: bool = False
_patched: bool = False
_lock: threading.Lock = threading.Lock()

# Number of batches to skip on the next enumerate(DataLoader) call.
# Set by _try_load_checkpoint() and consumed by the patched enumerate().
_dataloader_skip: int = 0

# ---------------------------------------------------------------------------
# Utility helpers
# ---------------------------------------------------------------------------


def _live_refs(ref_list: List[weakref.ref]) -> List[Any]:
    """Dereference weak-refs and return only live objects."""
    alive = []
    for r in ref_list:
        obj = r()
        if obj is not None:
            alive.append(obj)
    return alive


def _unwrap_model(model: Any) -> Any:
    """Unwrap ``DataParallel`` / ``DistributedDataParallel``."""
    while hasattr(model, "module"):
        model = model.module
    return model


def _has_parameters(module: Any) -> bool:
    """Return ``True`` if *module* owns at least one parameter."""
    try:
        next(module.parameters())
        return True
    except StopIteration:
        return False


def _root_models() -> List[Any]:
    """Return tracked models that are root-level and carry learnable parameters.

    Filters out sub-modules and stateless singletons (``nn.ReLU``, etc.).
    """
    all_models = _live_refs(_tracked_models)
    if not all_models:
        return []

    child_ids: Set[int] = set()
    for m in all_models:
        for child in m.modules():
            if child is not m:
                child_ids.add(id(child))

    return [
        m for m in all_models
        if id(m) not in child_ids and _has_parameters(m)
    ]


# ---------------------------------------------------------------------------
# RNG state helpers
# ---------------------------------------------------------------------------


def _collect_rng_states() -> Dict[str, Any]:
    """Snapshot all RNG states for deterministic resumption."""
    import torch

    rng: Dict[str, Any] = {}
    rng["torch_cpu"] = torch.random.get_rng_state()
    if torch.cuda.is_available():
        rng["torch_cuda"] = torch.cuda.get_rng_state_all()
    if "random" in sys.modules:
        import random
        rng["python"] = random.getstate()
    if "numpy" in sys.modules:
        import numpy
        rng["numpy"] = numpy.random.get_state()
    return rng


def _restore_rng_states(rng: Dict[str, Any]) -> None:
    """Restore RNG states captured by :func:`_collect_rng_states`."""
    import torch

    if "torch_cpu" in rng:
        torch.random.set_rng_state(rng["torch_cpu"])
    if "torch_cuda" in rng and torch.cuda.is_available():
        torch.cuda.set_rng_state_all(rng["torch_cuda"])
    if "python" in rng:
        import random
        random.setstate(rng["python"])
    if "numpy" in rng:
        import numpy
        numpy.random.set_state(rng["numpy"])


# ---------------------------------------------------------------------------
# Checkpoint save
# ---------------------------------------------------------------------------


def _save_checkpoint(path: str) -> None:
    """Serialize all tracked state to *path*."""
    import torch

    models = _root_models()
    optimizers = _live_refs(_tracked_optimizers)

    payload: Dict[str, Any] = {
        "models": [],
        "optimizers": [],
        "step_counts": [getattr(opt, "_hegm_step", 0) for opt in optimizers],
        "rng_states": _collect_rng_states(),
    }

    for idx, mdl in enumerate(models):
        inner = _unwrap_model(mdl)
        payload["models"].append({
            "idx": idx,
            "cls": type(inner).__qualname__,
            "state_dict": inner.state_dict(),
        })

    for idx, opt in enumerate(optimizers):
        payload["optimizers"].append({
            "idx": idx,
            "cls": type(opt).__qualname__,
            "state_dict": opt.state_dict(),
        })

    ckpt_dir = os.path.dirname(path)
    if ckpt_dir:
        os.makedirs(ckpt_dir, exist_ok=True)

    tmp_path = path + ".tmp"
    torch.save(payload, tmp_path)
    os.replace(tmp_path, path)


# ---------------------------------------------------------------------------
# Checkpoint load (auto-resume)
# ---------------------------------------------------------------------------


def _move_optimizer_state_to_device(optimizer: Any) -> None:
    """Move optimizer internal state tensors to match their parameter devices."""
    import torch

    for group in optimizer.param_groups:
        for p in group["params"]:
            if p not in optimizer.state:
                continue
            state = optimizer.state[p]
            for key, val in state.items():
                if isinstance(val, torch.Tensor) and val.device != p.device:
                    state[key] = val.to(p.device)


def _try_load_checkpoint() -> None:
    """Load checkpoint if it exists (called once on first optimizer creation)."""
    global _checkpoint_loaded

    with _lock:
        if _checkpoint_loaded:
            return
        _checkpoint_loaded = True

    if not os.path.isfile(CHECKPOINT_PATH):
        return

    import torch

    # Read checkpoint -------------------------------------------------------
    try:
        ckpt = torch.load(
            CHECKPOINT_PATH, map_location="cpu", weights_only=False
        )
    except TypeError:
        ckpt = torch.load(CHECKPOINT_PATH, map_location="cpu")
    except Exception as exc:
        warn(f"failed to read checkpoint: {exc}")
        return

    # Restore models --------------------------------------------------------
    models = _root_models()
    for entry in ckpt.get("models", []):
        idx = entry["idx"]
        if idx < len(models):
            inner = _unwrap_model(models[idx])
            try:
                inner.load_state_dict(entry["state_dict"])
                log(f"Restored model[{idx}] ({entry['cls']})")
            except Exception as exc:
                warn(f"model[{idx}] restore failed: {exc}")
        else:
            warn(f"checkpoint has model[{idx}] but only {len(models)} root model(s) tracked")

    # Restore optimizers ----------------------------------------------------
    optimizers = _live_refs(_tracked_optimizers)
    for entry in ckpt.get("optimizers", []):
        idx = entry["idx"]
        if idx < len(optimizers):
            try:
                optimizers[idx].load_state_dict(entry["state_dict"])
                _move_optimizer_state_to_device(optimizers[idx])
                log(f"Restored optimizer[{idx}] ({entry['cls']})")
            except Exception as exc:
                warn(f"optimizer[{idx}] restore failed: {exc}")
        else:
            warn(f"checkpoint has optimizer[{idx}] but only {len(optimizers)} optimizer(s) tracked")

    # Restore step counts ---------------------------------------------------
    global _dataloader_skip
    step_counts = ckpt.get("step_counts", [])
    for idx, count in enumerate(step_counts):
        if idx < len(optimizers):
            optimizers[idx]._hegm_step = count  # type: ignore[attr-defined]
            log(f"Optimizer[{idx}] resuming from global step {count}")

    # Tell the patched enumerate() how many batches to skip.
    if step_counts:
        _dataloader_skip = step_counts[0]

    # Restore RNG states ----------------------------------------------------
    rng = ckpt.get("rng_states")
    if rng:
        try:
            _restore_rng_states(rng)
            log("RNG states restored (torch/CUDA/Python/NumPy)")
        except Exception as exc:
            warn(f"RNG state restore failed: {exc}")

    log(f"Checkpoint resumed from {CHECKPOINT_PATH}")


# ---------------------------------------------------------------------------
# GPU resource release
# ---------------------------------------------------------------------------


def _release_gpu() -> None:
    """Best-effort CUDA teardown before exit."""
    import torch

    if not torch.cuda.is_available():
        return

    try:
        torch.cuda.synchronize()
    except Exception:
        pass

    for ref in _tracked_models:
        m = ref()
        if m is not None:
            try:
                m.cpu()
            except Exception:
                pass

    try:
        torch.cuda.empty_cache()
    except Exception:
        pass

    try:
        torch.cuda.synchronize()
    except Exception:
        pass

    log("GPU resources released.")


# ---------------------------------------------------------------------------
# Monkey-patching
# ---------------------------------------------------------------------------


def _apply_patches(_module: Any) -> None:
    """Wrap ``nn.Module.__init__`` and ``optim.Optimizer.__init__``."""
    global _patched

    with _lock:
        if _patched:
            return
        _patched = True

    import torch.nn as nn
    import torch.optim as optim

    # -- nn.Module.__init__ -------------------------------------------------
    _original_module_init = nn.Module.__init__

    @functools.wraps(_original_module_init)
    def _patched_module_init(self: nn.Module, *args: Any, **kwargs: Any) -> None:
        _original_module_init(self, *args, **kwargs)
        _tracked_models.append(weakref.ref(self))

    nn.Module.__init__ = _patched_module_init  # type: ignore[method-assign]

    # -- optim.Optimizer.__init__ -------------------------------------------
    _original_optimizer_init = optim.Optimizer.__init__

    @functools.wraps(_original_optimizer_init)
    def _patched_optimizer_init(
        self: optim.Optimizer, *args: Any, **kwargs: Any
    ) -> None:
        _original_optimizer_init(self, *args, **kwargs)
        _tracked_optimizers.append(weakref.ref(self))

        # Per-instance step() wrapper for step counting
        self._hegm_step = 0  # type: ignore[attr-defined]
        _orig_step = self.step

        def _counted_step(*a: Any, **kw: Any) -> Any:
            result = _orig_step(*a, **kw)
            self._hegm_step += 1  # type: ignore[attr-defined]
            return result

        self.step = _counted_step  # type: ignore[assignment]

        # Auto-resume on first optimizer creation
        _try_load_checkpoint()

    optim.Optimizer.__init__ = _patched_optimizer_init  # type: ignore[method-assign]

    log("Patches applied -- tracking torch.nn.Module & torch.optim.Optimizer")


# ---------------------------------------------------------------------------
# Backend class
# ---------------------------------------------------------------------------


class PyTorchBackend(AbstractBackend):
    """HeGM backend for PyTorch."""

    @property
    def name(self) -> str:
        return "pytorch"

    @property
    def target_module(self) -> str:
        return "torch"

    def apply_patches(self, module: Any) -> None:
        _apply_patches(module)

    def save_checkpoint(self, path: str) -> None:
        _save_checkpoint(path)

    def load_checkpoint(self, path: str) -> None:
        _try_load_checkpoint()

    def release_gpu(self) -> None:
        _release_gpu()

    def global_step(self, optimizer: Any = None) -> int:
        if optimizer is not None:
            return getattr(optimizer, "_hegm_step", 0)
        live = _live_refs(_tracked_optimizers)
        return getattr(live[0], "_hegm_step", 0) if live else 0

    def is_active(self) -> bool:
        return "torch" in sys.modules


# ---------------------------------------------------------------------------
# Auto-register
# ---------------------------------------------------------------------------

register(PyTorchBackend())
