"""Opportunistic Rust acceleration.

The Python/NumPy implementation is the REFERENCE and always works. If the
compiled wheels are installed (`pip install ./rust/cortexm-core` and
`pip install ./rust/quadrant`), the hot paths route through them:

* ``hashing.h64``          — keyed BLAKE2b (byte-exact parity, tested)
* ``vsa.bind/unbind``      — permutation gather with injected perms
* ``vsa.encode_fact``      — fused bind×3 + bundle + lexical mix + norm
* ``slb``                  — L1-resident lookaside buffer
* ``quadrant``             — page-clustered log-depth vector index

Design rule: the Rust side never GENERATES randomness — permutations and
role vectors are injected from Python's deterministic VSA state, so a
mixed Python/Rust deployment produces bit-identical holograms. Every
accelerated path has a NumPy fallback and a parity test
(``tests/test_rust_accel.py``).
"""

from __future__ import annotations

import os
from typing import Any

# --- optional wheels -------------------------------------------------------
try:                                            # pragma: no cover
    import cortexm_core as _core               # type: ignore
except Exception:                               # pragma: no cover
    _core = None

try:                                            # pragma: no cover
    import quadrant as _quadrant               # type: ignore
except Exception:                               # pragma: no cover
    _quadrant = None

RUST_AVAILABLE = _core is not None
QUADRANT_AVAILABLE = _quadrant is not None
_env = os.environ.get("CONTEXTM_RUST", "auto")
RUST_ENABLED = (_env == "1") or (_env == "auto" and RUST_AVAILABLE)
QUADRANT_ENABLED = (_env == "1") or (_env == "auto" and QUADRANT_AVAILABLE)


def rust_status() -> dict[str, Any]:
    return {
        "cortexm_core": bool(RUST_AVAILABLE and RUST_ENABLED),
        "quadrant": bool(QUADRANT_AVAILABLE and QUADRANT_ENABLED),
        "env": _env,
        "hint": ("" if RUST_AVAILABLE else
                 "build with: pip install ./rust/cortexm-core "
                 "./rust/quadrant (requires cargo + maturin)"),
    }


# --- hashing ---------------------------------------------------------------
if RUST_AVAILABLE and RUST_ENABLED:
    def h64(feature: str, seed: int = 0) -> int:
        return _core.h64(feature, seed)
else:                                           # pragma: no cover
    from cortexm.util import h64  # noqa: F401  (NumPy reference)


# --- VSA acceleration wrapper ----------------------------------------------
class RustVSA:
    """Binds a Python VSA to compiled bind/unbind/encode_fact.

    The authoritative perms live in the Python VSA; Rust receives them
    once via set_perm and accelerates the per-fact gathers afterwards.
    """

    def __init__(self, vsa) -> None:
        if not (RUST_AVAILABLE and RUST_ENABLED):
            raise RuntimeError("rust acceleration disabled or unavailable")
        self.vsa = vsa
        self._pb = _core.PermBindings(vsa.dims)
        self._injected: set[str] = set()

    def _inject(self, role: str) -> None:
        if role not in self._injected:
            self._pb.set_perm(role, self.vsa.perm(role).tolist())
            self._injected.add(role)

    def bind(self, role: str, filler):
        import numpy as np
        self._inject(role)
        f = np.ascontiguousarray(filler, dtype=np.float32)
        return self._pb.bind(role, f)

    def unbind(self, role: str, h):
        import numpy as np
        self._inject(role)
        hv = np.ascontiguousarray(h, dtype=np.float32)
        return self._pb.unbind(role, hv)

    def encode_fact(self, s_vec, r_vec, v_vec):
        import numpy as np
        for role in ("S", "R", "V"):
            self._inject(role)
        return self._pb.encode_fact(
            np.ascontiguousarray(s_vec, dtype=np.float32),
            np.ascontiguousarray(r_vec, dtype=np.float32),
            np.ascontiguousarray(v_vec, dtype=np.float32),
            self.vsa.lam)


# --- quadrant index wrapper --------------------------------------------------
class QuadrantANN:
    """Page-clustered log-depth index over holograms (approximate top-k).

    Falls back to exact numpy when the wheel is absent — callers get the
    same interface either way.
    """

    def __init__(self, vectors, page_capacity: int = 64) -> None:
        import numpy as np
        self.vectors = np.ascontiguousarray(vectors, dtype=np.float32)
        if QUADRANT_AVAILABLE and QUADRANT_ENABLED:
            self._idx = _quadrant.QuadrantIndex.build(
                self.vectors, page_capacity)
            self._exact = False
        else:                                    # pragma: no cover
            self._idx = None
            self._exact = True

    def search(self, query, k: int = 10, max_leaves: int = 16):
        import numpy as np
        q = np.ascontiguousarray(query, dtype=np.float32)
        if self._exact:                          # pragma: no cover
            sims = self.vectors @ q
            part = np.argpartition(-sims, k)[:k]
            order = part[np.argsort(-sims[part])]
            return order.tolist(), sims[order].tolist()
        ids, scores = self._idx.search(q, k, max_leaves)
        return list(ids), list(scores)

    def stats(self) -> dict:
        if self._exact:                          # pragma: no cover
            n = self.vectors.shape[0]
            return {"n_vectors": n, "mode": "exact-numpy",
                    "note": "quadrant wheel not installed"}
        return {"mode": "quadrant", "detail": self._idx.stats()}


# --- explicit binary/FP32 tiering ------------------------------------------
# User concern (architectural fix #1): "Binary HRR + FFT doesn't work →
# Separate tiers: binary for edge, FP32 for cloud."
# The codec stack already supports both (binary/RaBitQ for edge, PQ/INT8
# for cloud); this makes the routing EXPLICIT based on deployment tier.

EDGE_CODECS = ("binary", "rabitq")   # 96 B/v — fits on Raspberry Pi 5
CLOUD_CODECS = ("pq", "int8")        # 8 B/v or 770 B/v — bandwidth-dense


def detect_tier() -> str:
    """Auto-detect deployment tier from environment signals.

    Returns 'edge' or 'cloud'. Override via CONTEXTM_TIER env var.
    """
    explicit = os.environ.get("CONTEXTM_TIER", "").lower()
    if explicit in ("edge", "cloud"):
        return explicit
    # heuristic: cloud = multi-core + >4GB RAM + has scipy
    try:
        import multiprocessing
        if multiprocessing.cpu_count() >= 4:
            import os as _os
            mem = _os.sysconf("SC_PAGE_SIZE") * _os.sysconf("SC_PHYS_PAGES")
            if mem >= 4 * 1024**3:
                return "cloud"
    except Exception:
        pass
    return "edge"


def recommend_codec(tier: str | None = None, tmr: bool = True) -> str:
    """Recommend a codec for the deployment tier.

    Edge: binary (with TMR if self-healing required) — 96 B/v.
    Cloud: pq for bandwidth-dense, int8 for accuracy-critical.
    """
    t = tier or detect_tier()
    if t == "edge":
        return "binary" if tmr else "rabitq"
    return "pq"


def tier_status() -> dict:
    """Report current tier + recommended + active codec info."""
    tier = detect_tier()
    rec = recommend_codec(tier)
    return {
        "tier": tier,
        "recommended_codec": rec,
        "edge_codecs": list(EDGE_CODECS),
        "cloud_codecs": list(CLOUD_CODECS),
        "rust_enabled": RUST_ENABLED,
        "quadrant_enabled": QUADRANT_ENABLED,
        "note": (f"Auto-detected {tier} tier; recommend '{rec}' codec. "
                 f"Override with CONTEXTM_TIER=edge|cloud"),
    }


# --- SIMD kernel wrappers (Task 6-simd) ------------------------------------
# The Rust crate exposes `dot / dot_i8_f32 / cosine / l2_sq /
# batch_dot / batch_dot_i8 / topk / argmax` — each is a thin pyo3
# wrapper around the runtime-dispatched kernels in `rust/cortexm-core/
# src/simd.rs` (AVX-512 → AVX2+FMA → NEON → scalar). When the wheel is
# absent we fall back to numpy so callers always get a working answer.
#
# Design rule (mirrors RustVSA): the kernels NEVER allocate randomness
# and produce bit-compatible results across the Rust / NumPy paths (≤1e-5
# FP32 noise, asserted by `tests/test_rust_accel.py::TestSimdKernels`).
# Use the free-function form (`accel.cosine(a, b)`) for one-off queries;
# use the class form (`accel.SimdKernels().batch_dot(...)`) when you
# want to gate behaviour on `RUST_ENABLED` without re-checking globals.


def _as_f32(x) -> "np.ndarray":
    import numpy as np
    return np.ascontiguousarray(x, dtype=np.float32)


def dot(a, b) -> float:
    """SIMD dot product. NumPy fallback: `np.dot(a, b)`."""
    if RUST_AVAILABLE and RUST_ENABLED:
        import numpy as np
        return float(_core.dot(np.ascontiguousarray(a, dtype=np.float32),
                                np.ascontiguousarray(b, dtype=np.float32)))
    import numpy as np
    return float(np.dot(_as_f32(a), _as_f32(b)))


def dot_i8_f32(q8, q) -> float:
    """INT8 × f32 dot product. NumPy fallback casts int8 → f32 first."""
    if RUST_AVAILABLE and RUST_ENABLED:
        import numpy as np
        return float(_core.dot_i8_f32(
            np.ascontiguousarray(q8, dtype=np.int8),
            np.ascontiguousarray(q, dtype=np.float32)))
    import numpy as np
    q8 = np.ascontiguousarray(q8, dtype=np.int8)
    return float((q8.astype(np.float32) @ _as_f32(q)))


def cosine(a, b) -> float:
    """Cosine similarity via SIMD-accelerated dot + L2 norms.

    Identical vectors return exactly 1.0 (bit-exact through the
    SIMD self-dot path); distinct vectors agree with numpy's
    `a @ b / (|a|·|b|)` to within 1e-5 (FP32 lane-reduction noise).
    """
    if RUST_AVAILABLE and RUST_ENABLED:
        import numpy as np
        return float(_core.cosine(
            np.ascontiguousarray(a, dtype=np.float32),
            np.ascontiguousarray(b, dtype=np.float32)))
    import numpy as np
    a = _as_f32(a)
    b = _as_f32(b)
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12))


def l2_sq(a, b) -> float:
    """Squared L2 distance `sum((a-b)**2)`. 0.0 for identical vectors."""
    if RUST_AVAILABLE and RUST_ENABLED:
        import numpy as np
        return float(_core.l2_sq(
            np.ascontiguousarray(a, dtype=np.float32),
            np.ascontiguousarray(b, dtype=np.float32)))
    import numpy as np
    d = _as_f32(a) - _as_f32(b)
    return float(d @ d)


def batch_dot(rows, q, n_rows: int, dims: int) -> list[float]:
    """Matrix-vector product over a flat `[n_rows × dims]` f32 slice.

    Far more cache-friendly than calling `dot()` per row from Python —
    one boundary crossing for the whole batch.  Returns a plain list so
    callers can `np.asarray(...)` if they want a contiguous array.
    """
    if RUST_AVAILABLE and RUST_ENABLED:
        import numpy as np
        r = np.ascontiguousarray(rows, dtype=np.float32).reshape(-1)
        q = np.ascontiguousarray(q, dtype=np.float32)
        return list(_core.batch_dot(r, q, n_rows, dims))
    import numpy as np
    r = np.asarray(rows, dtype=np.float32).reshape(n_rows, dims)
    q = _as_f32(q)
    return (r @ q).tolist()


def batch_dot_i8(packed, q, n_rows: int, dims: int) -> list[float]:
    """INT8-packed rows × f32 query.  Raw int8·f32 dot products — callers
    apply per-row scales (the codec's aux array) themselves."""
    if RUST_AVAILABLE and RUST_ENABLED:
        import numpy as np
        p = np.ascontiguousarray(packed, dtype=np.int8).reshape(-1)
        q = np.ascontiguousarray(q, dtype=np.float32)
        return list(_core.batch_dot_i8(p, q, n_rows, dims))
    import numpy as np
    p = np.asarray(packed, dtype=np.int8).reshape(n_rows, dims)
    q = _as_f32(q)
    return (p.astype(np.float32) @ q).tolist()


def topk(scores, k: int) -> list[tuple[int, float]]:
    """Top-`k` (idx, score) tuples in descending order. O(N) via
    `select_nth_unstable`, then O(k log k) for the prefix sort."""
    if RUST_AVAILABLE and RUST_ENABLED:
        import numpy as np
        s = np.ascontiguousarray(scores, dtype=np.float32)
        return [(int(i), float(v)) for i, v in _core.topk(s, k)]
    import numpy as np
    s = np.asarray(scores, dtype=np.float32)
    k = max(0, min(k, s.size))
    if k == 0:
        return []
    part = np.argpartition(-s, k - 1)[:k]
    part = part[np.argsort(-s[part])]
    return [(int(i), float(s[i])) for i in part]


def argmax(scores) -> tuple[int, float]:
    """Return (idx, value) of the max element. Ties → first occurrence."""
    if RUST_AVAILABLE and RUST_ENABLED:
        import numpy as np
        s = np.ascontiguousarray(scores, dtype=np.float32)
        i, v = _core.argmax(s)
        return (int(i), float(v))
    import numpy as np
    s = np.asarray(scores, dtype=np.float32)
    if s.size == 0:
        return (0, 0.0)
    i = int(np.argmax(s))
    return (i, float(s[i]))


class SimdKernels:
    """Container class that mirrors the Rust kernel surface so callers
    can wire `accel.SimdKernels()` once and dispatch uniformly. The
    NumPy fallbacks produce bit-identical results within 1e-5 (FP32
    lane-reduction noise, asserted by parity tests).

    Mirrors the existing `RustVSA` pattern: the constructor raises
    `RuntimeError` if Rust is explicitly disabled (`CONTEXTM_RUST=0`)
    — when Rust is *available* but the user opted out, callers should
    use the free functions above instead.
    """

    def __init__(self) -> None:
        if not (RUST_AVAILABLE and RUST_ENABLED):
            raise RuntimeError(
                "SimdKernels requires Rust acceleration enabled "
                "(CONTEXTM_RUST=auto or =1, and the cortexm_core wheel built)")

    def dot(self, a, b) -> float:
        import numpy as np
        return float(_core.dot(
            np.ascontiguousarray(a, dtype=np.float32),
            np.ascontiguousarray(b, dtype=np.float32)))

    def dot_i8_f32(self, q8, q) -> float:
        import numpy as np
        return float(_core.dot_i8_f32(
            np.ascontiguousarray(q8, dtype=np.int8),
            np.ascontiguousarray(q, dtype=np.float32)))

    def cosine(self, a, b) -> float:
        import numpy as np
        return float(_core.cosine(
            np.ascontiguousarray(a, dtype=np.float32),
            np.ascontiguousarray(b, dtype=np.float32)))

    def l2_sq(self, a, b) -> float:
        import numpy as np
        return float(_core.l2_sq(
            np.ascontiguousarray(a, dtype=np.float32),
            np.ascontiguousarray(b, dtype=np.float32)))

    def batch_dot(self, rows, q, n_rows: int, dims: int) -> list[float]:
        import numpy as np
        r = np.ascontiguousarray(rows, dtype=np.float32).reshape(-1)
        q = np.ascontiguousarray(q, dtype=np.float32)
        return list(_core.batch_dot(r, q, n_rows, dims))

    def batch_dot_i8(self, packed, q, n_rows: int,
                      dims: int) -> list[float]:
        import numpy as np
        p = np.ascontiguousarray(packed, dtype=np.int8).reshape(-1)
        q = np.ascontiguousarray(q, dtype=np.float32)
        return list(_core.batch_dot_i8(p, q, n_rows, dims))

    def topk(self, scores, k: int) -> list[tuple[int, float]]:
        import numpy as np
        s = np.ascontiguousarray(scores, dtype=np.float32)
        return [(int(i), float(v)) for i, v in _core.topk(s, k)]

    def argmax(self, scores) -> tuple[int, float]:
        import numpy as np
        s = np.ascontiguousarray(scores, dtype=np.float32)
        i, v = _core.argmax(s)
        return (int(i), float(v))