"""Expert store: reads individual experts from the checkpoint file(s) with `pread`.

This is the SSD tier of the hierarchical memory. Every read is accounted (bytes, wall time) so that
stall time and SSD traffic can be reported per token. `nocache=True` sets F_NOCACHE on the file
descriptor (macOS) so that reads bypass the unified buffer cache for *new* pages; a cold measurement
still requires `purge` beforehand (see benchmarks/scripts/ssd_bandwidth.py).
"""
from __future__ import annotations

import fcntl
import os
import time
from typing import Dict, Tuple

import mlx.core as mx
import numpy as np

from concurrent.futures import ThreadPoolExecutor
from .registry import ModelSpec, TensorLoc

F_NOCACHE = 48
NP_DTYPE = {"F32": np.float32, "F16": np.float16, "U32": np.uint32, "I32": np.int32, "U8": np.uint8, "BF16": np.uint16}
MX_DTYPE = {"F32": mx.float32, "F16": mx.float16, "U32": mx.uint32, "I32": mx.int32, "U8": mx.uint8, "BF16": mx.bfloat16}


class ExpertStore:
    def __init__(self, spec: ModelSpec, nocache: bool = False, readers: int = 4):
        self.spec = spec
        self.nocache = nocache
        self._fds: Dict[str, int] = {}
        self.bytes_read = 0
        self.reads = 0
        self.read_time = 0.0
        self.readers = readers
        self._pool = ThreadPoolExecutor(max_workers=readers, thread_name_prefix="titan-store") if readers > 1 else None
        # an expert whose tensors are contiguous in the file is read with ONE pread (OLMoE: all 1 024 experts are contiguous)
        self._contig: Dict[Tuple[int, int], Tuple[str, int, int, list]] = {}
        for key, es in spec.experts.items():
            locs = sorted(es.tensors.items(), key=lambda kv: kv[1].offset)
            if all(a[1].file == b[1].file and a[1].offset + a[1].nbytes == b[1].offset for a, b in zip(locs, locs[1:])):
                self._contig[key] = (locs[0][1].file, locs[0][1].offset, sum(t.nbytes for _, t in locs), locs)

    def _fd(self, path: str) -> int:
        fd = self._fds.get(path)
        if fd is None:
            fd = os.open(path, os.O_RDONLY)
            if self.nocache:
                fcntl.fcntl(fd, F_NOCACHE, 1)
            self._fds[path] = fd
        return fd

    def read_tensor(self, loc: TensorLoc) -> mx.array:
        fd = self._fd(loc.file)
        t = time.perf_counter()
        buf = os.pread(fd, loc.nbytes, loc.offset)
        self.read_time += time.perf_counter() - t
        self.bytes_read += len(buf)
        self.reads += 1
        if len(buf) != loc.nbytes:
            raise IOError(f"short read for {loc.key}: {len(buf)} != {loc.nbytes}")
        arr = np.frombuffer(buf, dtype=NP_DTYPE[loc.dtype]).reshape(loc.shape)
        if loc.dtype == "BF16":
            return mx.array(arr).view(mx.bfloat16)
        return mx.array(arr)

    def _wrap(self, loc: TensorLoc, buf) -> mx.array:
        arr = np.frombuffer(buf, dtype=NP_DTYPE[loc.dtype]).reshape(loc.shape)
        if loc.dtype == "BF16":
            return mx.array(arr).view(mx.bfloat16)
        return mx.array(arr)

    def read_expert(self, layer: int, expert: int) -> Dict[str, mx.array]:
        c = self._contig.get((layer, expert))
        if c is None:
            es = self.spec.experts[(layer, expert)]
            return {name: self.read_tensor(loc) for name, loc in es.tensors.items()}
        path, off, n, locs = c
        fd = self._fd(path)
        t = time.perf_counter()
        buf = os.pread(fd, n, off)
        self.read_time += time.perf_counter() - t
        self.bytes_read += len(buf); self.reads += 1
        if len(buf) != n:
            raise IOError(f"short read for expert ({layer},{expert}): {len(buf)} != {n}")
        out = {}
        for name, loc in locs:
            a = loc.offset - off
            out[name] = self._wrap(loc, memoryview(buf)[a:a + loc.nbytes])
        return out

    def read_experts(self, layer: int, experts) -> Dict[int, Dict[str, mx.array]]:
        """Read several experts, in parallel when a reader pool exists (NVMe benefits from queue depth)."""
        experts = list(experts)
        if self._pool is None or len(experts) == 1:
            return {e: self.read_expert(layer, e) for e in experts}
        t = time.perf_counter()
        res = dict(zip(experts, self._pool.map(lambda e: self.read_expert(layer, e), experts)))
        self.read_time += 0  # per-read times already accumulated inside read_expert (sum of thread times, not wall)
        self.wall_batch_time = getattr(self, "wall_batch_time", 0.0) + (time.perf_counter() - t)
        return res

    def read_expert_bytes(self, layer: int, expert: int) -> Dict[str, bytes]:
        """Raw bytes (for the network transport)."""
        es = self.spec.experts[(layer, expert)]
        out = {}
        for name, loc in es.tensors.items():
            fd = self._fd(loc.file)
            t = time.perf_counter()
            out[name] = os.pread(fd, loc.nbytes, loc.offset)
            self.read_time += time.perf_counter() - t
            self.bytes_read += loc.nbytes
            self.reads += 1
        return out

    def stats(self) -> dict:
        return {"bytes_read": self.bytes_read, "reads": self.reads, "read_time_s": self.read_time}

    def reset_stats(self):
        self.bytes_read = 0
        self.reads = 0
        self.read_time = 0.0

    def close(self):
        for fd in self._fds.values():
            os.close(fd)
        self._fds.clear()