"""Dispatch-layer observe feed for the Bifrost provider (issue #17, Observe half).

`instrument_provider_app` wraps a built provider ASGI app so every inbound
bifrost-call emits one structured `OpEvent` — correlated by `session_id` read off
the dispatch JWT — WITHOUT touching the store's scope semantics (INV-004). It is
the lens that lets ratatoskr, owning BOTH ends of the round-trip, see exactly
which memory/affect ops a given turn produced.

The store-method stdout shim in `memory_store.py` cannot see `session_id` for
search/get/delete (bifrost withholds `ctx` from those store methods); this feed
sits at the DISPATCH/ASGI layer where the JWT — and thus `session_id` (its `sub`
claim) — is always present (INV-005).
"""
from __future__ import annotations

import base64
import json
import sys
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from typing import Any, Protocol


@dataclass(frozen=True)
class OpEvent:
    """One observed bifrost-call at the dispatch layer (scope-only, never content)."""

    ts: str  # ISO 8601 UTC, capture time
    plane: str  # "memory" | "affect"
    op: str  # verb: search / upsert_many / get / get_many / delete_many / emit / handshake
    session_id: str | None  # JWT `sub` at the dispatch layer; None only if the JWT omits it
    status: str  # "ok" | "error"
    req_summary: dict[str, Any]  # per-verb, scope-only — no record bodies
    resp_summary: dict[str, Any]  # per-verb counts + ids/scores — never verbatim content
    turn_id: str | None = None  # INV-005 reservation, literal; unused in v1


class OpSink(Protocol):
    """Anything that accepts an OpEvent. v1 ships JsonlOpSink; tests pass fakes."""

    def emit(self, event: OpEvent) -> None: ...


class JsonlOpSink:
    """Continuous append-only JSONL sink (INV-007: NOT per-session). One JSON line
    per OpEvent to a text stream (default stdout)."""

    def __init__(self, stream: Any = None) -> None:
        self._stream = sys.stdout if stream is None else stream

    def emit(self, event: OpEvent) -> None:
        self._stream.write(json.dumps(asdict(event), separators=(",", ":")) + "\n")
        self._stream.flush()


def maybe_instrument_from_env(app: Any, env: Any, *, plane: str) -> Any:
    """Opt-in serve wiring: when `RATATOSKR_OPFEED_PATH` is set, wrap `app` with
    the dispatch-layer op-feed writing JSONL to that path; otherwise return `app`
    unchanged. The append stream lives for the process (a long-running server)."""
    path = env.get("RATATOSKR_OPFEED_PATH")
    if not path:
        return app
    stream = open(path, "a", encoding="utf-8")  # process-lifetime append stream
    return instrument_provider_app(app, plane=plane, sink=JsonlOpSink(stream))


_BIFROST_PATHS = (
    "/bifrost/handshake",
    "/bifrost/memory-call",
    "/bifrost/affect-call",
)

_PLANE_BY_PATH = {
    "/bifrost/memory-call": "memory",
    "/bifrost/affect-call": "affect",
}


def _resolve_plane(configured: str, path: str) -> str:
    """For the combined provider (plane='combined', #18 D1) the OpEvent plane is
    derived from the request PATH — memory-call→memory, affect-call→affect,
    handshake→combined. A fixed plane ('memory'/'affect', the single-plane apps) is
    returned unchanged. The per-verb summary logic already keys on path, so only the
    plane STAMP changes."""
    if configured != "combined":
        return configured
    return _PLANE_BY_PATH.get(path, "combined")


def _b64url_decode(seg: str) -> bytes:
    return base64.urlsafe_b64decode(seg + "=" * (-len(seg) % 4))


def _session_id_from_auth(auth: bytes | None) -> str | None:
    """Read the `sub` claim (= session_id, per bifrost DispatchContext) off the
    dispatch JWT WITHOUT verifying its signature — the inner app does real
    verification; we only read a claim for correlation. None if absent/malformed."""
    if not auth:
        return None
    try:
        token = auth.decode("latin-1").strip()
        if token.lower().startswith("bearer "):
            token = token[7:].strip()
        parts = token.split(".")
        if len(parts) != 3:
            return None
        payload = json.loads(_b64url_decode(parts[1]))
        sub = payload.get("sub")
        return sub if isinstance(sub, str) else None
    except Exception:
        return None


def _op_from(path: str, req: dict[str, Any]) -> str:
    """The verb: 'handshake' for the handshake path; otherwise the body's
    `operation`, with the affect-plane `affect.` prefix stripped (affect.emit ->
    emit) so op vocabulary stays bare per the contract."""
    if path == "/bifrost/handshake":
        return "handshake"
    operation = req.get("operation") or "unknown"
    if path == "/bifrost/affect-call" and operation.startswith("affect."):
        return operation.split(".", 1)[1]
    return operation


def _ids_summary(args: dict[str, Any]) -> list[Any]:
    """Mirror bifrost `_ids_arg`: ids | chunk_ids | [chunk_id|id]."""
    ids = args.get("ids") or args.get("chunk_ids")
    if ids is None:
        single = args.get("chunk_id") or args.get("id")
        ids = [single] if single is not None else []
    return ids


def _req_summary(plane: str, path: str, op: str, req: dict[str, Any]) -> dict[str, Any]:
    """Scope-only request summary — NEVER record bodies / PAD content."""
    if path == "/bifrost/handshake":
        # The handshake REQUEST field is `capabilities` (bifrost reference_server
        # _protocol.py:181 reads request_body["capabilities"]) — NOT the transposed
        # `capabilities_requested`, which never existed on the wire (caps_requested
        # was silently always None). Fixed per the heid-code-review #17 catch.
        return {"caps_requested": req.get("capabilities")}
    if plane == "affect":
        return {}  # affect stays conduit-opaque — no PAD content surfaced
    args = req.get("args") or {}
    if op == "search":
        return {
            "scope_all": args.get("scope_all") or {},
            "scope_any": args.get("scope_any") or [],
            "top_k": args.get("top_k"),
        }
    if op == "upsert_many":
        records = args.get("records") or []
        return {
            "record_count": len(records),
            "scopes": [r.get("scope") for r in records],
        }
    if op in ("get", "get_many", "delete_many"):
        return {"ids": _ids_summary(args)}
    return {}


def _resp_summary(
    plane: str, path: str, op: str, resp: dict[str, Any], status: str
) -> dict[str, Any]:
    """Per-verb counts + ids/scores — never verbatim content. On error, the
    bifrost error `code` (INV-007: failures recorded, not hidden)."""
    if status == "error":
        return {"error": resp.get("code")}
    if path == "/bifrost/handshake":
        return {"ok": True, "caps_granted": resp.get("capabilities_granted")}
    if plane == "affect":
        return {"stored": bool(resp.get("stored"))}
    if op == "search":
        results = resp.get("results") or []
        return {
            "hit_count": len(results),
            "hits": [
                {"chunk_id": r.get("chunk_id"), "score": r.get("score")}
                for r in results
            ],
        }
    if op == "upsert_many":
        return {"upserted": resp.get("upserted"), "replayed": resp.get("replayed")}
    if op == "get":
        return {"found_count": 1 if resp.get("record") else 0}
    if op == "get_many":
        return {"found_count": len(resp.get("records") or [])}
    if op == "delete_many":
        return {"deleted": resp.get("deleted")}
    return {}


def _build_event(
    plane: str, path: str, scope: dict[str, Any], req_body: bytes, captured: dict[str, Any]
) -> OpEvent:
    plane = _resolve_plane(plane, path)  # 'combined' → per-path; fixed plane unchanged
    headers = dict(scope.get("headers") or [])
    session_id = _session_id_from_auth(headers.get(b"authorization"))
    status = "ok" if 200 <= int(captured["status"]) < 300 else "error"
    req = _safe_json(req_body)
    resp = _safe_json(captured["body"])

    op = _op_from(path, req)
    return OpEvent(
        ts=datetime.now(UTC).isoformat(),
        plane=plane,
        op=op,
        session_id=session_id,
        status=status,
        req_summary=_req_summary(plane, path, op, req),
        resp_summary=_resp_summary(plane, path, op, resp, status),
    )


def _safe_json(raw: bytes) -> dict[str, Any]:
    if not raw:
        return {}
    try:
        value = json.loads(raw)
    except (json.JSONDecodeError, ValueError):
        return {}
    return value if isinstance(value, dict) else {}


def instrument_provider_app(app: Any, *, plane: str, sink: OpSink) -> Any:
    """Wrap a built provider ASGI `app` so each inbound bifrost-call emits one
    OpEvent to `sink`. Read-only over dispatch — store scope semantics untouched
    (INV-004). A sink/summary failure never propagates into the dispatch path
    (POST-003 / INV-007) — it is swallowed and logged to stderr.
    """
    if plane not in ("memory", "affect", "combined"):
        raise ValueError(
            f"plane must be 'memory', 'affect', or 'combined', got {plane!r}"
        )

    async def wrapped(scope: dict[str, Any], receive: Any, send: Any) -> None:
        if scope.get("type") != "http" or scope.get("path") not in _BIFROST_PATHS:
            await app(scope, receive, send)
            return

        # Buffer the request body so we can BOTH summarise it AND replay it to the
        # inner app (the inner app consumes `receive`; we must not steal the body).
        req_body = b""
        more = True
        while more:
            message = await receive()
            if message["type"] == "http.request":
                req_body += message.get("body", b"")
                more = message.get("more_body", False)
            else:  # http.disconnect
                more = False

        replayed = False

        async def replay_receive() -> dict[str, Any]:
            nonlocal replayed
            if not replayed:
                replayed = True
                return {"type": "http.request", "body": req_body, "more_body": False}
            return {"type": "http.disconnect"}

        captured: dict[str, Any] = {"status": 500, "body": b""}

        async def capture_send(message: dict[str, Any]) -> None:
            if message["type"] == "http.response.start":
                captured["status"] = message["status"]
            elif message["type"] == "http.response.body":
                captured["body"] += message.get("body", b"")
            await send(message)

        await app(scope, replay_receive, capture_send)

        try:
            sink.emit(_build_event(plane, scope["path"], scope, req_body, captured))
        except Exception as exc:  # observe gap, never a serve break (POST-003)
            sys.stderr.write(f"[opfeed] OpEvent emit failed (swallowed): {exc!r}\n")

    return wrapped
