"""Content-addressable hyperdimensional storage.

A cell is one hypervector of ``d`` bits, identified by a *name*.  The address
of a name is itself a hypervector (``hvc("cell/" + name)``), so the store has
no linear address bus: you hand it a vector and it answers with the cell whose
address vector it most resembles.  Exact-name access is served by a dictionary
index to keep the VM deterministic; near-miss access falls through to a
similarity scan, which is the behaviour that makes it content-addressable.

Two structures live here:

``cells``      the data store — program instructions, register shadows, tool
               records, structure records, all the same kind of object
``cleanup``    a symbol table mapping hypervectors back to names, which is what
               turns a noisy decoded field into an opcode / register / literal

The file format (``.hdc``) is deliberately boring and byte-stable: a header,
the packed bit-vectors, then the symbol table.  Vectors are regenerated from
names where they are symbolic, so an image written on one machine loads on
another at the same ``d``.
"""

from __future__ import annotations

import os
import struct
from typing import Iterable, Iterator

from .hypervec import DEFAULT_DIM, hvc, mask, popcount, sim

MAGIC = b"HDCV1\x00"
#: A fuzzy address match must beat this to resolve without an exact hit.
FUZZY_FLOOR = 0.20


def _nbytes(d: int) -> int:
    return (d + 7) // 8


def pack(hv: int, d: int = DEFAULT_DIM) -> bytes:
    """Pack a hypervector into fixed-width little-endian bytes."""
    return (hv & mask(d)).to_bytes(_nbytes(d), "little")


def unpack(raw: bytes, d: int = DEFAULT_DIM) -> int:
    """Inverse of :func:`pack`."""
    return int.from_bytes(raw, "little") & mask(d)


class HDCStorage:
    """A hyperdimensional store: named cells plus a cleanup symbol table."""

    def __init__(self, d: int = DEFAULT_DIM, capacity: int = 8192, label: str = "hdc0"):
        self.d = d
        self.capacity = capacity
        self.label = label
        self.cells: dict[str, int] = {}
        self.order: list[str] = []
        self._by_addr: dict[int, str] = {}
        self.cleanup: dict[int, str] = {}
        self._cleanup_names: list[str] = []
        #: Monotonic counters, surfaced by ``stat`` and the IDE.
        self.reads = 0
        self.writes = 0
        self.fuzzy_reads = 0

    # ---------------------------------------------------------------- cells
    def addr(self, name: str) -> int:
        """The address hypervector of a cell name."""
        return hvc("cell/" + name, self.d)

    def write(self, name: str, hv: int, symbol: bool = False) -> None:
        """Store ``hv`` at ``name``.  Returns nothing; overwrites in place."""
        if name not in self.cells:
            if len(self.order) >= self.capacity:
                raise MemoryError(
                    "HDC storage full (%d cells)" % self.capacity
                )
            self.order.append(name)
        self.cells[name] = hv & mask(self.d)
        self._by_addr[self.addr(name)] = name
        self.writes += 1
        if symbol:
            self.define(name)

    def read(self, name: str, default: int | None = None) -> int:
        """Read a cell by name."""
        if name in self.cells:
            self.reads += 1
            return self.cells[name]
        if default is not None:
            return default
        raise KeyError("no such HDC cell: %r" % (name,))

    def remove(self, name: str) -> bool:
        """Delete a cell and its symbol entry; False when it was not present.

        A filesystem has to be able to take cells away when a file is deleted,
        otherwise ``len(storage)`` grows forever and every stat lies about it.
        """
        if name not in self.cells:
            return False
        del self.cells[name]
        try:
            self.order.remove(name)
        except ValueError:
            pass
        address = self.addr(name)
        self._by_addr.pop(address, None)
        self.cleanup.pop(address, None)
        try:
            self._cleanup_names.remove(name)
        except ValueError:
            pass
        return True

    def define(self, name: str) -> int:
        """Add a name to cleanup memory; returns its address vector."""
        hv = self.addr(name)
        if hv not in self.cleanup:
            self.cleanup[hv] = name
            self._cleanup_names.append(name)
        return hv

    def __contains__(self, name: object) -> bool:
        return name in self.cells

    def __len__(self) -> int:
        return len(self.cells)

    def __iter__(self) -> Iterator[tuple[str, int]]:
        for name in self.order:
            yield name, self.cells[name]

    # -------------------------------------------------- content addressing
    def resolve(self, addr_hv: int, floor: float = FUZZY_FLOOR) -> tuple[str, float] | None:
        """Map an address vector to a cell name.

        Exact hit first (confidence 1.0).  Otherwise the nearest cell address
        by similarity, provided it clears ``floor``.  This is the read path
        that gives the machine its HDC character: a corrupted address still
        lands on the right cell and reports how sure it was.
        """
        name = self._by_addr.get(addr_hv & mask(self.d))
        if name is not None:
            self.reads += 1
            return name, 1.0
        best_name: str | None = None
        best = -2.0
        for a, n in self._by_addr.items():
            s = sim(a, addr_hv, self.d)
            if s > best:
                best, best_name = s, n
        if best_name is not None and best >= floor:
            self.reads += 1
            self.fuzzy_reads += 1
            return best_name, best
        return None

    def recall(
        self, probe: int, k: int = 5, prefix: str = "", floor: float = 0.0
    ) -> list[tuple[str, int, float]]:
        """Associative search: the ``k`` cell *values* most similar to ``probe``.

        This is the compute-in-memory primitive — no index, no address, the
        search happens by comparing against every stored vector.
        """
        scored: list[tuple[str, int, float]] = []
        for name in self.order:
            if prefix and not name.startswith(prefix):
                continue
            hv = self.cells[name]
            s = sim(hv, probe, self.d)
            if s >= floor:
                scored.append((name, hv, s))
        scored.sort(key=lambda t: t[2], reverse=True)
        return scored[:k]

    def symbols(self) -> list[str]:
        """Cleanup memory contents, in registration order."""
        return list(self._cleanup_names)

    def nearest_symbol(self, hv: int, k: int = 1, prefix: str = "") -> list[tuple[str, float]]:
        """Cleanup: nearest named symbol(s) for a noisy hypervector."""
        scored: list[tuple[str, float]] = []
        for name in self._cleanup_names:
            if prefix and not name.startswith(prefix):
                continue
            scored.append((name, sim(hvc("cell/" + name, self.d), hv, self.d)))
        scored.sort(key=lambda t: t[1], reverse=True)
        return scored[:k]

    # -------------------------------------------------------------- chex
    def bitplanes(self, name: str, width: int = 64) -> list[str]:
        """Bit planes of a cell, WSB first — the IDE's raw view."""
        from .hypervec import bitplanes

        return bitplanes(self.read(name), self.d, width)

    def stat(self) -> dict:
        """Storage statistics (cheap; used by ``MEMSTAT`` and the IDE)."""
        used = len(self.order)
        total_bits = used * self.d
        return {
            "label": self.label,
            "d": self.d,
            "cells_used": used,
            "cells_free": max(0, self.capacity - used),
            "bytes": used * (_nbytes(self.d) + 24),
            "symbols": len(self.cleanup),
            "reads": self.reads,
            "writes": self.writes,
            "fuzzy_reads": self.fuzzy_reads,
            "occupancy": used / self.capacity if self.capacity else 0.0,
            "stored_bits": total_bits,
            "total_popcount": sum(popcount(v) for v in self.cells.values()),
        }

    # --------------------------------------------------------- persistence
    def save(self, path: str) -> int:
        """Write the whole store as a ``.hdc`` image.  Returns bytes written."""
        nb = _nbytes(self.d)
        with open(path, "wb") as fh:
            fh.write(MAGIC)
            fh.write(struct.pack("<IIII", self.d, self.capacity, len(self.order), len(self._cleanup_names)))
            for name in self.order:
                raw = name.encode("utf-8")
                fh.write(struct.pack("<H", len(raw)))
                fh.write(raw)
                fh.write(pack(self.cells[name], self.d))
            for name in self._cleanup_names:
                raw = name.encode("utf-8")
                fh.write(struct.pack("<H", len(raw)))
                fh.write(raw)
        return os.path.getsize(path)

    @classmethod
    def load(cls, path: str, label: str | None = None) -> "HDCStorage":
        """Load a ``.hdc`` image."""
        with open(path, "rb") as fh:
            if fh.read(len(MAGIC)) != MAGIC:
                raise ValueError("not an HDC image: %s" % path)
            d, capacity, n_cells, n_syms = struct.unpack("<IIII", fh.read(16))
            store = cls(d=d, capacity=max(capacity, n_cells), label=label or os.path.basename(path))
            nb = _nbytes(d)
            for _ in range(n_cells):
                (ln,) = struct.unpack("<H", fh.read(2))
                name = fh.read(ln).decode("utf-8")
                hv = unpack(fh.read(nb), d)
                store.write(name, hv)
            for _ in range(n_syms):
                (ln,) = struct.unpack("<H", fh.read(2))
                store.define(fh.read(ln).decode("utf-8"))
        return store

    # ------------------------------------------------------------- helpers
    def clone(self) -> "HDCStorage":
        other = HDCStorage(d=self.d, capacity=self.capacity, label=self.label)
        for name in self.order:
            other.write(name, self.cells[name])
        for name in self._cleanup_names:
            other.define(name)
        return other

    def dump(self, limit: int = 24, width: int = 48) -> str:
        """Human-readable listing of the store."""
        lines = [
            "HDC storage %r  d=%d  cells=%d/%d  symbols=%d"
            % (self.label, self.d, len(self.order), self.capacity, len(self.cleanup))
        ]
        for i, name in enumerate(self.order[:limit]):
            hv = self.cells[name]
            from .hypervec import bits

            b = bits(hv, self.d)
            lines.append(
                "  %4d  %-22s %s...  pop=%d"
                % (i, name, b[:width], popcount(hv))
            )
        if len(self.order) > limit:
            lines.append("  ... %d more cells" % (len(self.order) - limit))
        return "\n".join(lines)


def storage_from_pairs(pairs: Iterable[tuple[str, int]], d: int = DEFAULT_DIM) -> HDCStorage:
    """Build a store from ``(name, hypervector)`` pairs."""
    items = list(pairs)
    store = HDCStorage(d=d, capacity=max(64, len(items)))
    for name, hv in items:
        store.write(name, hv)
    return store