"""MemoryStore: tensor-granular storage for Phase 1 and Phase 2 (ADR-0020 D7).

Logically byte-addressable, implemented as addr → numpy ndarray mapping.
Read/write are reference-based (no copy) for Phase 1 performance.
"""
from __future__ import annotations

import numpy as np


# numpy dtype string → numpy dtype mapping
_DTYPE_MAP = {
    "f16": np.float16,
    "f32": np.float32,
    "f64": np.float64,
    "bf16": np.float16,  # numpy has no bfloat16; use float16 as proxy
    "i8": np.int8,
    "i16": np.int16,
    "i32": np.int32,
    "i64": np.int64,
    "u8": np.uint8,
    "u16": np.uint16,
    "u32": np.uint32,
}


def _resolve_dtype(dtype: str) -> np.dtype:
    if dtype in _DTYPE_MAP:
        return np.dtype(_DTYPE_MAP[dtype])
    return np.dtype(dtype)


class MemoryStore:
    """Tensor-granular memory storage (ADR-0020 D7).

    Stores numpy ndarrays by (space, addr) key.
    Write = reference store (no copy), read = reference return (no copy).
    Overwrite at same addr replaces the entire tensor.
    """

    def __init__(self) -> None:
        # {space: {addr: ndarray}}
        self._storage: dict[str, dict[int, np.ndarray]] = {}

    def write(self, space: str, addr: int, data: np.ndarray) -> None:
        """Store tensor at (space, addr). Reference-only, no copy."""
        if space not in self._storage:
            self._storage[space] = {}
        self._storage[space][addr] = data

    def read(self, space: str, addr: int, shape: tuple[int, ...] | None = None,
             dtype: str | None = None) -> np.ndarray:
        """Read tensor from (space, addr). Returns reference, no copy.

        Models real Triton ``tl.load(ptr + offset, shape=...)`` semantics:

          - exact base hit, exact bytes: returns the stored array (optionally
            reinterpreted/reshaped — byte-conserving).
          - exact base hit, fewer bytes: returns a prefix slice.
          - non-base addr falling inside a stored region: returns the
            slice at ``[addr - base, addr - base + requested_nbytes)``.

        Out-of-bounds, no-containing-region, and dtype mismatches still
        raise. Used by tile-sweep kernels that load sub-slices at
        ``ptr + tile_start * row_bytes``.
        """
        store = self._storage.get(space)
        if store is None:
            raise KeyError(f"No data at ({space}, 0x{addr:x})")

        arr = store.get(addr)
        base_addr = addr
        if arr is None:
            base_addr, arr = self._find_containing_region(store, addr)
            if arr is None:
                raise KeyError(f"No data at ({space}, 0x{addr:x})")

        if dtype is not None:
            np_dtype = _resolve_dtype(dtype)
            if arr.dtype != np_dtype:
                arr = arr.view(np_dtype)

        if shape is None:
            if base_addr != addr:
                raise ValueError(
                    f"Partial read at ({space}, 0x{addr:x}) requires explicit shape"
                )
            return arr

        requested_nbytes = int(np.prod(shape)) * arr.dtype.itemsize
        byte_offset = addr - base_addr
        if byte_offset + requested_nbytes > arr.nbytes:
            raise ValueError(
                f"Out-of-bounds read at ({space}, 0x{addr:x}): "
                f"offset {byte_offset}B + {requested_nbytes}B > "
                f"region {arr.nbytes}B"
            )

        if byte_offset == 0 and requested_nbytes == arr.nbytes:
            return arr if arr.shape == shape else arr.reshape(shape)

        itemsize = arr.dtype.itemsize
        if byte_offset % itemsize != 0 or requested_nbytes % itemsize != 0:
            raise ValueError(
                f"Unaligned partial read at ({space}, 0x{addr:x}): "
                f"offset {byte_offset}B / size {requested_nbytes}B "
                f"not a multiple of {itemsize}B"
            )
        flat = arr.reshape(-1)
        start = byte_offset // itemsize
        end = start + requested_nbytes // itemsize
        return flat[start:end].reshape(shape)

    @staticmethod
    def _find_containing_region(
        store: dict[int, np.ndarray], addr: int,
    ) -> tuple[int, np.ndarray | None]:
        """Find the stored region whose ``[base, base+nbytes)`` contains ``addr``.

        Linear scan over stored bases. Acceptable since the store is
        small (one entry per deployed tensor); upgrade to a sorted
        interval map only if profiling shows it dominates.
        """
        for base, arr in store.items():
            if base <= addr < base + arr.nbytes:
                return base, arr
        return 0, None

    def has(self, space: str, addr: int) -> bool:
        return addr in self._storage.get(space, {})

    def snapshot(self) -> MemoryStore:
        """Create a shallow copy for Phase 2 initialization."""
        new = MemoryStore()
        for space, addrs in self._storage.items():
            new._storage[space] = dict(addrs)  # shallow copy of addr→ndarray map
        return new
