"""Asynchronous prefetcher: a small pool of reader threads fetches predicted experts from the ExpertStore (SSD tier)
while the GPU computes the current layer; completed fetches are installed into the slot table on the main thread when the
layer is reached (`drain`). Predictors come from titan.predictors and only see information available at decision time.

Accounting per layer: prefetch_issued, prefetch_installed, prefetch_used (installed experts that were subsequently hit),
prefetch_wasted (installed but evicted before use), prefetch_late (needed before the fetch completed: counted as a demand
miss whose stall is the *remaining* latency).
"""
from __future__ import annotations

import threading, time
from concurrent.futures import Future, ThreadPoolExecutor
from typing import Dict, List, Optional

import mlx.core as mx
import numpy as np

from .slots import SlotTable
from .store import ExpertStore


class Prefetcher:
    def __init__(self, tables: List[SlotTable], store: ExpertStore, predictor, m: int = 8, workers: int = 2):
        self.tables = tables
        self.store = store
        self.pred = predictor
        self.m = m
        self.pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="titan-prefetch")
        self.inflight: List[Dict[int, Future]] = [dict() for _ in tables]
        self.lock = threading.Lock()
        self.issued = [0] * len(tables); self.installed = [0] * len(tables); self.late = [0] * len(tables); self.dup = [0] * len(tables)
        self.prev_topk: Optional[np.ndarray] = None
        self.prev_probs: Optional[np.ndarray] = None
        self.cur: Dict = {}
        self.L = len(tables)
        self.K = store.spec.top_k

    # ---- reading ---------------------------------------------------------------
    def _fetch(self, layer: int, expert: int):
        t = time.perf_counter()
        tensors = self.store.read_expert(layer, expert)  # pread + np -> mx (CPU side, no GPU work)
        return tensors, time.perf_counter() - t

    def issue(self, layer: int, experts: List[int]):
        tab = self.tables[layer]
        for e in experts[: self.m]:
            if e in tab.slot_of or e in self.inflight[layer]:
                self.dup[layer] += 1
                continue
            self.inflight[layer][e] = self.pool.submit(self._fetch, layer, e)
            self.issued[layer] += 1

    def drain(self, layer: int, protected: List[int], wait_for: List[int] = ()):
        """Install completed prefetches for `layer`; optionally wait for the ones in `wait_for` (needed now)."""
        tab = self.tables[layer]
        inf = self.inflight[layer]
        for e in list(inf):
            fut = inf[e]
            if e in wait_for and not fut.done():
                t = time.perf_counter()
                fut.result()
                tab.stall_time += time.perf_counter() - t  # remaining latency of a late prefetch
                self.late[layer] += 1
            if fut.done():
                tensors, _ = fut.result()
                del inf[e]
                if tab.insert_prefetched(e, tensors, protected) is not None:
                    self.installed[layer] += 1

    # ---- hooks called by the execution block ---------------------------------------
    def before_layer(self, layer: int, needed: List[int]):
        """Called with the materialised routing of `layer` before its experts are ensured."""
        self.drain(layer, needed, wait_for=[e for e in needed if e in self.inflight[layer]])
        if hasattr(self.pred, "set_true"):
            pass

    def after_layer(self, layer: int, uniq: List[int], probs: Optional[np.ndarray] = None):
        self.cur[layer] = np.array(uniq)
        if probs is not None:
            self.cur[("probs", layer)] = probs
        if self.pred.horizon == "layer" and layer + 1 < self.L:
            self.issue(layer + 1, [int(e) for e in self.pred.predict(layer + 1, self.prev_topk, self.prev_probs, self.cur, self.m)])
        if layer == self.L - 1:
            self.prev_topk = np.array([self.cur.get(l, np.array([]))[: self.K] if len(self.cur.get(l, [])) >= self.K else np.pad(self.cur.get(l, np.array([])), (0, self.K - len(self.cur.get(l, [])))) for l in range(self.L)])
            self.prev_probs = np.stack([self.cur[("probs", l)] for l in range(self.L)]) if all(("probs", l) in self.cur for l in range(self.L)) else None
            self.cur = {}
            if self.pred.horizon == "token":
                for l in range(self.L):
                    self.issue(l, [int(e) for e in self.pred.predict(l, self.prev_topk, self.prev_probs, {}, self.m)])

    def stats(self) -> dict:
        return {"issued": sum(self.issued), "installed": sum(self.installed), "late": sum(self.late), "dup": sum(self.dup),
                "predictor": self.pred.name, "m": self.m}

    def reset_stats(self):
        self.issued = [0] * self.L; self.installed = [0] * self.L; self.late = [0] * self.L; self.dup = [0] * self.L

    def close(self):
        self.pool.shutdown(wait=True)