"""Spawn one D3 process per GPU so FLA/Triton mutable state is not shared."""

import multiprocessing
import threading
import traceback
from concurrent.futures import ProcessPoolExecutor

from ._compat import Engine, NativeAbstention, Unsupported

TARGET = f"{__package__}.engine:D3Engine"
RUN_SEED = 20260919

_ENGINE = None
_FAILURE = None
_START_LOCK = threading.Lock()


class WorkerError(RuntimeError):
    def __init__(self, record):
        super().__init__(record["error"])
        self.exception_name = record["exception"]
        self.remote_traceback = record["traceback"]


def _invoke(target, options, method, args, kwargs):
    global _ENGINE, _FAILURE
    if _FAILURE is not None:
        return False, _FAILURE
    try:
        if _ENGINE is None:
            from importlib import import_module

            if target == TARGET:
                import torch

                torch.manual_seed(RUN_SEED)
            module, name = target.split(":")
            _ENGINE = getattr(import_module(module), name)(**options)
        if method == "info":
            value = (_ENGINE.provenance, _ENGINE.runtime(), _ENGINE.latency)
        else:
            value = getattr(_ENGINE, method)(*args, **kwargs)
            if method == "__call__":
                _ENGINE.synchronize()
        return True, value
    except Exception as exc:
        record = {"error": str(exc), "exception": type(exc).__name__,
                  "traceback": traceback.format_exc()}
        if isinstance(exc, NativeAbstention):
            record.update(kind="abstained", raw=exc.raw)
        elif isinstance(exc, Unsupported):
            record["kind"] = "unsupported"
        if _ENGINE is None or record["exception"] in ("OutOfMemoryError", "AcceleratorError") or "device-side assert" in str(exc):
            _FAILURE = record
        return False, record


class _Child:
    def __init__(self, options, target=TARGET):
        self.options = options
        self.target = target
        self.lock = threading.Lock()
        self.pool = ProcessPoolExecutor(max_workers=1, mp_context=multiprocessing.get_context("spawn"))
        try:
            self.info = self.call("info")
        except BaseException:
            self.pool.shutdown(wait=True, cancel_futures=True)
            raise

    def call(self, method, *args, **kwargs):
        with self.lock:
            ok, value = self.pool.submit(_invoke, self.target, self.options, method, args, kwargs).result()
        if ok:
            return value
        if value.get("kind") == "unsupported":
            raise Unsupported(value["error"])
        if value.get("kind") == "abstained":
            raise NativeAbstention(value["raw"])
        raise WorkerError(value)

    def close(self):
        # Process exit releases the model and all CUDA allocations, even after an error.
        self.pool.shutdown(wait=True, cancel_futures=True)


class D3ProcessEngine(Engine):
    name = "d3"

    def __init__(self, device=None, replicas=None, **options):
        import torch

        gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
        if replicas is not None and int(replicas) > 1:
            if int(replicas) > gpus:
                raise ValueError(f"replicas={replicas} needs more than {gpus} CUDA devices")
            self.devices = [f"cuda:{i}" for i in range(int(replicas))]
        elif device in (None, "cuda"):
            self.devices = [f"cuda:{i}" for i in range(gpus)] or ["cpu"]
        else:
            self.devices = [str(device)]
        super().__init__(device=device, **options)
        self._options = options
        self._children = {}
        self._local = threading.local()
        self._lock = threading.Lock()
        self._next = 0
        child = self._child(0)
        provenance, runtime, self.latency = child.info
        self.provenance = {**provenance, "gpu_execution": "one spawned process per GPU"}
        self._runtime = runtime

    def _child(self, slot):
        child = self._children.get(slot)
        if child is not None:
            return child
        with self._lock:
            if slot not in self._children:
                # Avoid simultaneous checkpoint loads exhausting host RAM.
                with _START_LOCK:
                    self._children[slot] = _Child({**self._options, "device": self.devices[slot]})
            return self._children[slot]

    def __call__(self, *args, **kwargs):
        slot = getattr(self._local, "slot", None)
        if slot is None:
            with self._lock:
                slot = self._next
                self._next = (self._next + 1) % len(self.devices)
            self._local.slot = slot
        return self._child(slot).call("__call__", *args, **kwargs)

    def warmup(self):
        self._child(0).call("warmup")

    def synchronize(self):
        for child in list(self._children.values()):
            child.call("synchronize")

    def runtime(self):
        return {**self._runtime, "gpu_execution": "spawn", "devices": self.devices}

    def close(self):
        for child in self._children.values():
            child.close()
        self._children.clear()