"""SQLite + sqlite-vec durable memory store (Bifrost memory plane, v1 basic plane).

Contract: docs/contracts/bifrost_memory_provider.contract.md

Worldtree writes Tier-3 agent memory chunks here (`upsert_many`) and recalls them
by vector similarity (`search`), with point reads (`get`/`get_many`) and deletes
(`delete_many`). We persist each chunk verbatim and read only its structural
surface — embedding (rank), scope (isolation), id + revision (optimistic lock),
origin/injection_source (consistency rule). Semantic content is never interpreted.
We implement bifrost's OWN MemoryDataStore Protocol; conformance is #195 parity
vs InMemoryMemoryStore.
"""
from __future__ import annotations

import hashlib
import json
import logging
import sqlite3
import sys
import time
from typing import Any

import sqlite_vec
from bifrost.consumer import ConsumerRegistration, build_memory_app
from bifrost.memory import (
    IdempotencyConflict,
    InvalidArguments,
    InvalidFilter,
    RevisionMismatch,
    StoreCapabilities,
)
from bifrost.reference_server import JwtVerifier

_SHORT_RETRY_TTL_SECONDS = 300
_DURABLE_JOB_TTL_SECONDS = 24 * 60 * 60

# bifrost wire v0.6 scope lattice: three subject axes + agent_self (the #248
# agent-identity axis, canonical since #10/v0.5). An axis outside it is InvalidFilter
# (-> memory.invalid_filter 400), matching bifrost's reference _validate_scope.
_SCOPE_LATTICE = {"end_user", "group", "tenant", "agent_self"}

# Inbound memory-call observability (#17 observe brick). A self-contained
# stdout handler so the lines reliably reach the provider's stdout regardless
# of uvicorn's logging config. INFO-level, no propagation to root.
_log = logging.getLogger("ratatoskr.provider.memory")
if not _log.handlers:
    _h = logging.StreamHandler(sys.stdout)
    _h.setFormatter(logging.Formatter("%(asctime)s [memory-provider] %(message)s"))
    _log.addHandler(_h)
    _log.setLevel(logging.INFO)
    _log.propagate = False


def _ctx_actor(ctx: Any) -> str:
    """Reference `_ctx_actor`: actor = job_id | jwt_sub | session_id (never the record)."""
    return str(
        getattr(ctx, "job_id", None)
        or getattr(ctx, "jwt_sub", None)
        or getattr(ctx, "session_id", "")
    )


def _payload_digest(value: Any) -> str:
    """Reference digest: sha256 of canonical JSON (sorted keys, compact, str-coerced)."""
    blob = json.dumps(value, sort_keys=True, separators=(",", ":"), default=str)
    return hashlib.sha256(blob.encode("utf-8")).hexdigest()


def _idempotency_ttl_seconds(idempotency_class: str | None) -> int:
    if idempotency_class == "durable-job":
        return _DURABLE_JOB_TTL_SECONDS
    return _SHORT_RETRY_TTL_SECONDS


def _chunk_id(record: dict) -> str:
    """Reference `_chunk_id`: first non-empty of id / chunk_id / memory_id."""
    for key in ("id", "chunk_id", "memory_id"):
        value = record.get(key)
        if isinstance(value, str) and value:
            return value
    raise InvalidArguments("record missing id")


def _record_vector(record: dict) -> list[float]:
    """Reference `_record_vector`: embedding, falling back to vector, else []."""
    value = record.get("embedding", record.get("vector", []))
    return [float(v) for v in value] if isinstance(value, list) else []


def _scope_subset(record_scope: dict, filter_dict: dict) -> bool:
    """True iff record_scope has EVERY axis of filter_dict (AND). Empty filter = match."""
    return all(record_scope.get(axis) == value for axis, value in filter_dict.items())


def _matches_scope(record_scope: Any, scope_all: dict, scope_any: list) -> bool:
    """INV-005 (wire v0.6, #11): a record passes iff
    `(scope_all empty OR record ⊇ scope_all) AND (scope_any empty OR it matches ≥1 element)`.
    scope_any is OR/union over a LIST of conjunctive scope dicts (each AND-matched as a whole),
    never single axes flattened together (the Worldtree #297 over-broadening foot-gun).
    Byte-faithful to bifrost reference `reference_server/memory.py:_matches_scope`.
    """
    scope = record_scope if isinstance(record_scope, dict) else {}
    if not _scope_subset(scope, scope_all):
        return False
    if scope_any and not any(_scope_subset(scope, element) for element in scope_any):
        return False
    return True


def _validate_scope(scope_all: dict, scope_any: list) -> None:
    """STEP 1: scope_all is a flat dict, scope_any a list of flat dicts; every axis in BOTH
    fields ∈ the v0.6 lattice {end_user, group, tenant, agent_self} (else InvalidFilter),
    matching the reference `_validate_scope`."""
    if not isinstance(scope_all, dict):
        raise InvalidArguments("scope_all must be a flat {axis: value} dict")
    if not isinstance(scope_any, list) or any(not isinstance(e, dict) for e in scope_any):
        raise InvalidArguments("scope_any must be a list of {axis: value} dicts")
    for scope in (scope_all, *scope_any):
        if any(axis not in _SCOPE_LATTICE for axis in scope):
            raise InvalidFilter("scope_filter contains unsupported axis")


def _validate_injection(record: dict) -> None:
    """INV-007: injected_context requires injection_source; non-injected forbids it."""
    origin = record.get("origin")
    injection_source = record.get("injection_source")
    if origin == "injected_context" and not injection_source:
        raise InvalidArguments("injected_context record requires injection_source")
    if injection_source and origin != "injected_context":
        raise InvalidArguments("injection_source only valid for injected_context origin")


class RatatoskrMemoryStore:
    """The MemoryDataStore-shaped store handed to bifrost's build_memory_app."""

    def __init__(self, conn: sqlite3.Connection, embedding_dim: int):
        self._conn = conn
        self._dim = embedding_dim

    def describe_store(self) -> dict:
        # INV-006: advertise ONLY the v1 basic-plane capabilities (advertise-=>-implement).
        return StoreCapabilities(
            relational_edges_supported=False,
            optimistic_locking_supported=True,
            atomic_supersede_supported=False,
            transaction_supported=False,
            filterable_metadata_fields=[],
        ).to_dict()

    async def upsert_many(
        self,
        records: list[dict],
        *,
        idempotency_key: str,
        ctx: Any,
        expected_revisions: dict | None = None,
        idempotency_class: str | None = None,
    ) -> dict:
        if not (isinstance(idempotency_key, str) and idempotency_key):  # PRE-001
            raise InvalidArguments("idempotency_key required")
        _log.info(
            "memory-call upsert_many REQUEST: %d record(s) idempotency_key=%s actor=%s scopes=%s",
            len(records), idempotency_key, _ctx_actor(ctx),
            [r.get("scope") for r in records],
        )
        # INV-002: idempotency_id = ("default", verb, actor-from-ctx, key); digest over payload.
        digest = _payload_digest({"records": records, "expected_revisions": expected_revisions})
        idempotency_id = "|".join(("default", "upsert_many", _ctx_actor(ctx), idempotency_key))
        cached = self._conn.execute(
            "SELECT digest FROM memory_idempotency WHERE idempotency_id = ?",
            (idempotency_id,),
        ).fetchone()
        if cached is not None:
            if cached[0] == digest:
                return {"upserted": len(records), "replayed": True}
            raise IdempotencyConflict("idempotency key reused with different payload")
        for record in records:  # INV-007 / PRE-002: validate before any write
            _validate_injection(record)
        # INV-004: all chunk rows + vec rows + the idempotency record in one transaction;
        # a RevisionMismatch raised here rolls the whole batch back.
        with self._conn:
            if expected_revisions:  # INV-003: optimistic lock, checked before any write
                for record in records:
                    cid = _chunk_id(record)
                    if cid in expected_revisions:
                        row = self._conn.execute(
                            "SELECT revision FROM memory_chunks WHERE chunk_id = ?", (cid,)
                        ).fetchone()
                        current = row[0] if row is not None else 0
                        if current != expected_revisions[cid]:
                            raise RevisionMismatch(
                                f"stale expected revision for {cid}: "
                                f"{expected_revisions[cid]} != {current}"
                            )
            for record in records:
                chunk_id = _chunk_id(record)
                # INV-003: first insert -> revision 1; re-upsert -> revision + 1.
                self._conn.execute(
                    "INSERT INTO memory_chunks "
                    "(chunk_id, record_json, revision, scope_json, origin) "
                    "VALUES (?, ?, 1, ?, ?) "
                    "ON CONFLICT(chunk_id) DO UPDATE SET "
                    "record_json=excluded.record_json, revision=memory_chunks.revision + 1, "
                    "scope_json=excluded.scope_json, origin=excluded.origin",
                    (
                        chunk_id,
                        json.dumps(record),
                        json.dumps(record.get("scope")),
                        record.get("origin"),
                    ),
                )
                self._conn.execute("DELETE FROM memory_vec WHERE chunk_id = ?", (chunk_id,))
                self._conn.execute(
                    "INSERT INTO memory_vec(chunk_id, embedding) VALUES (?, ?)",
                    (chunk_id, sqlite_vec.serialize_float32(_record_vector(record))),
                )
            self._conn.execute(
                "INSERT INTO memory_idempotency (idempotency_id, digest, expires_at) "
                "VALUES (?, ?, ?)",
                (idempotency_id, digest, time.time() + _idempotency_ttl_seconds(idempotency_class)),
            )
        return {"upserted": len(records), "replayed": False}

    async def search(
        self,
        vector: list[float],
        *,
        top_k: int,
        scope_all: dict | None = None,
        scope_any: list | None = None,
        metadata_filter: dict | None = None,
        include: dict | None = None,
        fidelity_target: Any = None,
    ) -> list[dict]:
        if len(vector) != self._dim:  # PRE-001
            raise InvalidArguments(f"vector length {len(vector)} != embedding_dim {self._dim}")
        if metadata_filter:  # PRE-002: v1 advertises no filterable metadata fields
            raise InvalidArguments("metadata_filter is unsupported in v1")
        scope_all = scope_all or {}
        scope_any = scope_any or []
        _validate_scope(scope_all, scope_any)  # STEP 1 (raises InvalidArguments / InvalidFilter)
        _log.info(
            "memory-call search REQUEST: scope_all=%r scope_any=%r top_k=%s vec_dim=%d",
            scope_all, scope_any, top_k, len(vector),
        )

        def _emit(rs: list[dict]) -> list[dict]:
            _log.info(
                "memory-call search RESPONSE: %d hit(s) %s",
                len(rs),
                [
                    {
                        "chunk_id": r["chunk_id"],
                        "score": round(r["score"], 4),
                        "scope": r["chunk"].get("scope"),
                    }
                    for r in rs
                ],
            )
            return rs

        if top_k <= 0:  # POST-001: at most top_k
            return _emit([])
        total = self._conn.execute("SELECT COUNT(*) FROM memory_vec").fetchone()[0]
        if total == 0:
            return _emit([])
        # Over-fetch every candidate ranked by cosine distance, then scope-filter and
        # take top_k — so top_k counts IN-SCOPE hits (INV-005), not pre-filter hits.
        rows = self._conn.execute(
            "SELECT v.chunk_id, v.distance, c.record_json, c.revision "
            "FROM memory_vec v JOIN memory_chunks c ON c.chunk_id = v.chunk_id "
            f"WHERE v.embedding MATCH ? AND k = {total} ORDER BY v.distance",
            (sqlite_vec.serialize_float32([float(x) for x in vector]),),
        ).fetchall()
        results: list[dict] = []
        for chunk_id, distance, record_json, revision in rows:
            record = json.loads(record_json)
            if not _matches_scope(record.get("scope"), scope_all, scope_any):
                continue
            results.append(
                {
                    "chunk": record,
                    "chunk_id": chunk_id,
                    "score": 1.0 - distance,  # vec0 cosine distance -> similarity
                    "recalled_view": record.get("distillate", record),
                    "revision": revision,
                }
            )
            if len(results) >= top_k:
                break
        return _emit(results)

    async def get(self, chunk_id: str) -> dict | None:
        # INV-001: verbatim round-trip + an attached revision key, or None.
        row = self._conn.execute(
            "SELECT record_json, revision FROM memory_chunks WHERE chunk_id = ?",
            (chunk_id,),
        ).fetchone()
        if row is None:
            return None
        record = json.loads(row[0])
        record["revision"] = row[1]
        return record

    async def get_many(self, ids: list[str]) -> list[dict]:
        # List form of get: found records only (absent ids are skipped).
        found = []
        for chunk_id in ids:
            record = await self.get(chunk_id)
            if record is not None:
                found.append(record)
        return found

    async def delete_many(self, ids: list[str]) -> dict:
        # One transaction: chunk row + its vec row leave together (no orphan vec rows).
        _log.info("memory-call delete_many REQUEST: ids=%s", ids)
        deleted = 0
        with self._conn:
            for chunk_id in ids:
                cur = self._conn.execute(
                    "DELETE FROM memory_chunks WHERE chunk_id = ?", (chunk_id,)
                )
                if cur.rowcount > 0:
                    deleted += 1
                self._conn.execute("DELETE FROM memory_vec WHERE chunk_id = ?", (chunk_id,))
        return {"deleted": deleted}


def open_memory_store(db_path: str, *, embedding_dim: int) -> RatatoskrMemoryStore:
    """Open the SQLite+sqlite-vec memory store, creating schema + the vec index on first use."""
    if not (isinstance(embedding_dim, int) and embedding_dim > 0):  # PRE-002
        raise ValueError("embedding_dim must be a positive int")
    # check_same_thread=False: the memory provider is an ASGI app; uvicorn/Starlette
    # (and TestClient always) may run a handler off the connection's creating thread.
    # The event loop serializes the sync sqlite calls, so this is safe. Mirrors the
    # affect store (bifrost_affect_provider INV-006); surfaced by a TestClient-driven
    # memory-call search through the combined provider (#18 D1).
    conn = sqlite3.connect(db_path, check_same_thread=False)
    conn.enable_load_extension(True)
    sqlite_vec.load(conn)
    conn.enable_load_extension(False)
    conn.execute("PRAGMA busy_timeout=5000")  # wait up to 5s, don't fail SQLITE_BUSY at once
    if db_path != ":memory:":
        conn.execute("PRAGMA journal_mode=WAL")
    conn.execute(
        "CREATE TABLE IF NOT EXISTS memory_chunks ("
        "chunk_id TEXT PRIMARY KEY, record_json TEXT NOT NULL, "
        "revision INTEGER NOT NULL, scope_json TEXT, origin TEXT)"
    )
    conn.execute(
        "CREATE TABLE IF NOT EXISTS memory_idempotency ("
        "idempotency_id TEXT PRIMARY KEY, digest TEXT NOT NULL, expires_at REAL)"
    )
    # INV-005 / cosine recall: vec0 index at the fixed PINNED_EMBEDDER_DIM, cosine metric.
    conn.execute(
        "CREATE VIRTUAL TABLE IF NOT EXISTS memory_vec USING vec0("
        f"chunk_id TEXT PRIMARY KEY, embedding float[{embedding_dim}] distance_metric=cosine)"
    )
    conn.commit()
    return RatatoskrMemoryStore(conn, embedding_dim)


def build_memory_provider_app(
    store: RatatoskrMemoryStore,
    heimdall_key: bytes,
    consumer_id: str = "ratatoskr",
):
    """Wire the JWT verifier + registration and hand the store to bifrost.

    Returns a Starlette ASGI app exposing POST /bifrost/handshake and
    POST /bifrost/memory-call. The library owns the wire; this is the thin glue.
    """
    if not isinstance(store.describe_store(), dict):  # PRE-001 / INV-008
        raise ValueError("store must advertise capabilities via describe_store()")
    if not (isinstance(heimdall_key, bytes) and heimdall_key):  # PRE-002
        raise ValueError("heimdall_key must be non-empty bytes")
    verifier = JwtVerifier(algorithm="HS256", key_bytes=heimdall_key)
    registration = ConsumerRegistration(consumer_id=consumer_id)
    return build_memory_app(store=store, verifier=verifier, registration=registration)
