from __future__ import annotations

import multiprocessing as mp
import queue
import time
from typing import Dict, Optional

import torch

from hippotorch.memory.consolidator import Consolidator
from hippotorch.memory.store import MemoryStore


class AsyncConsolidationWorker:
    """Run consolidation in a background process to avoid blocking training."""

    def __init__(
        self,
        consolidator: Consolidator,
        memory: MemoryStore,
        *,
        start_method: str = "spawn",
        device: torch.device | str | None = None,
    ) -> None:
        self.consolidator = consolidator
        self.memory = memory
        self.ctx = mp.get_context(start_method)
        self._requests: mp.Queue = self.ctx.Queue()
        self._results: mp.Queue = self.ctx.Queue()
        self._process: mp.Process | None = None
        self._device = (
            torch.device(device)
            if device is not None
            else next(consolidator.parameters()).device
        )
        self._job_counter = 0
        self._pending: Dict[int, float] = {}
        self._stats: Dict[str, float] = {
            "pending_jobs": 0,
            "queue_depth": 0,
            "last_duration": 0.0,
            "last_swap_duration": 0.0,
            "last_applied": 0.0,
            "dropped_results": 0.0,
        }

    def start(self) -> None:
        if self._process is not None and self._process.is_alive():
            return
        worker_consolidator = self.consolidator.cpu()
        self._process = self.ctx.Process(
            target=_worker_loop,
            args=(worker_consolidator, self._requests, self._results),
            daemon=True,
        )
        self._process.start()
        self.consolidator.to(self._device)

    def submit(
        self,
        *,
        steps: int = 1,
        batch_size: int = 32,
        refresh_keys: bool = True,
        report_quality: bool = False,
        reward_weighted: bool = True,
    ) -> None:
        """Queue a consolidation job using a CPU snapshot of the memory store."""

        self.start()
        snapshot = self.memory.clone(device="cpu")
        metadata = snapshot.snapshot_metadata()
        self._job_counter += 1
        job_id = self._job_counter
        self._pending[job_id] = time.perf_counter()
        try:
            queue_depth = self._requests.qsize()
        except (NotImplementedError, AttributeError):
            queue_depth = -1
        if queue_depth >= 0:
            self._stats["queue_depth"] = float(queue_depth)
        self._stats["pending_jobs"] = float(len(self._pending))
        self._requests.put(
            {
                "job_id": job_id,
                "memory": snapshot,
                "metadata": metadata,
                "steps": steps,
                "batch_size": batch_size,
                "refresh_keys": refresh_keys,
                "report_quality": report_quality,
                "reward_weighted": reward_weighted,
            }
        )

    def poll(
        self, *, apply: bool = True, timeout: float = 0.0
    ) -> Optional[Dict[str, float]]:
        """Return metrics if a background job completed."""

        try:
            result = self._results.get(timeout=timeout)
        except queue.Empty:
            return None

        job_id = result.get("job_id")
        if job_id is not None:
            start = self._pending.pop(job_id, None)
            if start is not None:
                self._stats["last_duration"] = time.perf_counter() - start
            self._stats["pending_jobs"] = float(len(self._pending))

        if apply:
            self._apply_result(result)
        return result.get("metrics")

    def shutdown(self) -> None:
        if self._process is None:
            return
        self._requests.put({"stop": True})
        self._process.join(timeout=1.0)
        if self._process.is_alive():  # pragma: no cover - safety net
            self._process.terminate()
        self._process = None
        self._pending.clear()

    def _apply_result(self, payload: Dict) -> None:
        encoder_state = payload.get("encoder_state")
        if encoder_state:
            self.consolidator.encoder.load_state_dict(encoder_state)
        optimizer_state = payload.get("optimizer_state")
        if optimizer_state:
            self.consolidator.optimizer.load_state_dict(optimizer_state)
        self.consolidator.total_steps = int(
            payload.get("steps", self.consolidator.total_steps)
        )

        key_tensor: torch.Tensor | None = payload.get("keys")
        metadata = payload.get("metadata") or {}
        if key_tensor is not None and metadata:
            swap_start = time.perf_counter()
            applied, dropped = self.memory.apply_key_snapshot(
                keys=key_tensor,
                key_timestamps=payload.get(
                    "key_timestamps", self.memory.key_timestamps
                ),
                episode_ids=metadata.get("episode_ids", []),
                slot_versions=metadata.get("slot_versions", []),
            )
            self._stats["last_swap_duration"] = time.perf_counter() - swap_start
            self._stats["last_applied"] = float(applied)
            self._stats["dropped_results"] += float(dropped)

        self.consolidator.to(self._device)

    def stats(self) -> Dict[str, float]:
        """Return a snapshot of async worker diagnostics."""

        return dict(self._stats)


def _worker_loop(
    consolidator: Consolidator, request_q: mp.Queue, result_q: mp.Queue
) -> None:
    consolidator = consolidator.cpu()
    while True:
        task = request_q.get()
        if isinstance(task, dict) and task.get("stop"):
            break
        memory: MemoryStore = task["memory"]
        metadata = task.get("metadata", {})
        metrics = consolidator.sleep(
            memory,
            steps=task.get("steps", 1),
            batch_size=task.get("batch_size", 32),
            refresh_keys=task.get("refresh_keys", True),
            report_quality=task.get("report_quality", False),
            reward_weighted=task.get("reward_weighted", True),
        )
        payload = {
            "job_id": task.get("job_id"),
            "metrics": metrics,
            "encoder_state": consolidator.encoder.state_dict(),
            "optimizer_state": consolidator.optimizer.state_dict(),
            "steps": consolidator.total_steps,
            "keys": memory.keys.detach(),
            "key_timestamps": list(memory.key_timestamps),
            "metadata": metadata,
        }
        result_q.put(payload)


__all__ = ["AsyncConsolidationWorker"]
