"""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.

        If shape/dtype match stored tensor, returns as-is.
        If dtype differs, performs reinterpret cast (view).
        If shape differs but nbytes match, reshapes.
        """
        store = self._storage.get(space)
        if store is None or addr not in store:
            raise KeyError(f"No data at ({space}, 0x{addr:x})")
        arr = store[addr]
        if dtype is not None:
            np_dtype = _resolve_dtype(dtype)
            if arr.dtype != np_dtype:
                arr = arr.view(np_dtype)
        if shape is not None and arr.shape != shape:
            if arr.nbytes != np.prod(shape) * arr.dtype.itemsize:
                raise ValueError(
                    f"Shape mismatch: stored {arr.shape} ({arr.nbytes}B) "
                    f"vs requested {shape} ({np.prod(shape) * arr.dtype.itemsize}B)"
                )
            arr = arr.reshape(shape)
        return arr

    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
