"""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 .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):
        self.spec = spec
        self.nocache = nocache
        self._fds: Dict[str, int] = {}
        self.bytes_read = 0
        self.reads = 0
        self.read_time = 0.0

    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 read_expert(self, layer: int, expert: int) -> Dict[str, mx.array]:
        es = self.spec.experts[(layer, expert)]
        return {name: self.read_tensor(loc) for name, loc in es.tensors.items()}

    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()