"""Multi-modality witness for canonical-shape questions (Ticket #000028).

When a question matches a canonical shape (arithmetic@v1,
logic-kernel@v1, future kernels), fan out three independent answer
paths in parallel:

- **kernel**: deterministic π* — ground truth.
- **cache**: prior commitment from ``providence_cache`` (closure
  returns ``None`` when no prior row, e.g. pre-Ticket #000027).
- **alien (LLM)**: an independently-trained model invited to answer
  the same question.

Each modality's answer is canonicalized through the kernel for
byte-comparison. Three paths byte-equal → ``STRICT-WITNESSED`` —
the strongest grounding signal arborist can offer (mono-modal ×
independent-training × prior-commitment). Disagreement is recorded
as falsification signal, not an error.

Default OFF — tripling LLM cost on every arithmetic question is a
tax operators don't want by default. Enable per-call via
``--witness`` CLI flag or per-policy via
``policy["canonical_witness_enabled"] = True``.
"""

from __future__ import annotations

import time
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutTimeoutError
from dataclasses import dataclass, field
from typing import Callable

from arborist.qa.client import ChatClient


WITNESS_SYSTEM_PROMPT = (
    "You are a deterministic math/logic computer. Reply with ONLY the "
    "canonical answer — no prose, no explanation, no LaTeX, no leading "
    '"=", no trailing punctuation. For arithmetic emit a fraction (e.g. '
    "3/10) or integer. For propositional logic emit TRUE or FALSE. "
    "Output nothing else."
)


@dataclass(frozen=True)
class ModalityResult:
    """One modality's contribution to the witness."""

    modality: str                       # "kernel" | "cache" | "llm"
    raw_answer: str | None              # what the modality said (display)
    canonical_bytes: bytes | None       # post-canonicalization bytes (None on reject)
    error: str | None                   # "PIS_REJECT" / "TIMEOUT" / "ABSENT" / etc.
    elapsed_ms: float

    @property
    def ok(self) -> bool:
        return self.canonical_bytes is not None and self.error is None

    def to_dict(self) -> dict:
        return {
            "modality": self.modality,
            "raw_answer": self.raw_answer,
            "canonical_bytes_hex": (
                self.canonical_bytes.hex() if self.canonical_bytes else None
            ),
            "error": self.error,
            "elapsed_ms": round(self.elapsed_ms, 1),
            "ok": self.ok,
        }


@dataclass(frozen=True)
class Witness:
    """Outcome of a multi-modality witness for one canonical-shape question."""

    pi_star_ref: str
    canonical_answer_bytes: bytes
    modalities: dict[str, ModalityResult] = field(default_factory=dict)
    agreement_label: str = "UNCLASSIFIED"
    elapsed_ms: float = 0.0

    def to_dict(self) -> dict:
        return {
            "pi_star_ref": self.pi_star_ref,
            "canonical_answer_bytes_hex": self.canonical_answer_bytes.hex(),
            "canonical_answer_text": self.canonical_answer_bytes.decode(
                "utf-8", errors="replace"
            ),
            "modalities": {
                name: m.to_dict() for name, m in self.modalities.items()
            },
            "agreement_label": self.agreement_label,
            "elapsed_ms": round(self.elapsed_ms, 1),
        }


def _canonicalize_via_kernel(text: str | None, pi_star_ref: str) -> bytes | None:
    """Run the named kernel on ``text``. Returns canonical bytes, or
    ``None`` when the kernel rejects (PiStarError) or input is empty.
    """
    from arborist.pi_star import PiStarError, get
    if text is None:
        return None
    cleaned = text.strip()
    if not cleaned:
        return None
    try:
        return get(pi_star_ref).canonicalize(cleaned.encode("utf-8"))
    except PiStarError:
        return None
    except Exception:  # pragma: no cover — defensive
        return None


def _classify_agreement(
    modalities: dict[str, ModalityResult],
    canonical_answer_bytes: bytes,
) -> str:
    """Map (kernel, cache, llm) outcomes → agreement label.

    The kernel is the ground-truth reference; cache & llm are checked
    *against* the kernel's canonical_answer_bytes.
    """
    cache = modalities.get("cache")
    llm = modalities.get("llm")

    cache_present = cache is not None and cache.ok
    llm_present = llm is not None and llm.ok

    cache_agrees = (
        cache_present and cache.canonical_bytes == canonical_answer_bytes
    )
    llm_agrees = (
        llm_present and llm.canonical_bytes == canonical_answer_bytes
    )

    cache_disagrees = cache_present and not cache_agrees
    llm_disagrees = llm_present and not llm_agrees

    # Decision tree (kernel always present when this is called):
    if cache_agrees and llm_agrees:
        return "STRICT-WITNESSED"
    if cache_agrees and llm_disagrees:
        return "LLM-DIVERGED"
    if cache_disagrees and llm_agrees:
        return "CACHE-DRIFT"
    if cache_disagrees and llm_disagrees:
        return "LLM-AND-CACHE-DIVERGED"
    if not cache_present and llm_agrees:
        return "KERNEL-LLM-AGREE"
    if not cache_present and llm_disagrees:
        return "KERNEL-LLM-DIVERGED"
    if cache_agrees and not llm_present:
        return "KERNEL-CACHE-AGREE"
    if cache_disagrees and not llm_present:
        return "CACHE-DRIFT-LLM-ABSENT"
    # Both cache and llm absent — only the kernel ran. Degenerate but
    # surface it explicitly so the audit-line tail makes sense.
    return "KERNEL-ONLY"


def _run_kernel_modality(
    canonical_answer_bytes: bytes, pi_star_ref: str
) -> ModalityResult:
    """Kernel modality: by definition produces canonical_answer_bytes.

    The work was already done before run_witness was called — we
    receive the answer and just package it as a ModalityResult so
    the matrix has a kernel row alongside cache + llm.
    """
    return ModalityResult(
        modality="kernel",
        raw_answer=canonical_answer_bytes.decode("utf-8", errors="replace"),
        canonical_bytes=canonical_answer_bytes,
        error=None,
        elapsed_ms=0.0,
    )


def _run_cache_modality(
    cache_lookup: Callable[[], bytes | None] | None,
    pi_star_ref: str,
) -> ModalityResult:
    """Cache modality: invoke the lookup closure. ``None`` → ABSENT
    (no prior commitment); bytes → canonicalize & report.
    """
    t0 = time.monotonic()
    if cache_lookup is None:
        return ModalityResult(
            modality="cache",
            raw_answer=None,
            canonical_bytes=None,
            error="ABSENT",
            elapsed_ms=0.0,
        )
    try:
        raw_bytes = cache_lookup()
    except Exception as e:  # pragma: no cover
        return ModalityResult(
            modality="cache",
            raw_answer=None,
            canonical_bytes=None,
            error=f"LOOKUP_ERROR:{type(e).__name__}",
            elapsed_ms=(time.monotonic() - t0) * 1000.0,
        )
    elapsed = (time.monotonic() - t0) * 1000.0
    if raw_bytes is None:
        return ModalityResult(
            modality="cache",
            raw_answer=None,
            canonical_bytes=None,
            error="ABSENT",
            elapsed_ms=elapsed,
        )
    raw_text = raw_bytes.decode("utf-8", errors="replace")
    canonical = _canonicalize_via_kernel(raw_text, pi_star_ref)
    if canonical is None:
        return ModalityResult(
            modality="cache",
            raw_answer=raw_text,
            canonical_bytes=None,
            error="PIS_REJECT",
            elapsed_ms=elapsed,
        )
    return ModalityResult(
        modality="cache",
        raw_answer=raw_text,
        canonical_bytes=canonical,
        error=None,
        elapsed_ms=elapsed,
    )


def _run_llm_modality(
    question: str,
    chat_client: ChatClient | None,
    model_id: str,
    pi_star_ref: str,
) -> ModalityResult:
    """LLM modality: ask the alien to compute the answer with strict
    output formatting; canonicalize the response through the kernel.
    """
    t0 = time.monotonic()
    if chat_client is None:
        return ModalityResult(
            modality="llm",
            raw_answer=None,
            canonical_bytes=None,
            error="ABSENT",
            elapsed_ms=0.0,
        )
    messages = [
        {"role": "system", "content": WITNESS_SYSTEM_PROMPT},
        {"role": "user", "content": question},
    ]
    try:
        raw = chat_client.chat_completion(
            messages,
            model=model_id,
            temperature=0.0,
            max_tokens=64,
            top_p=1.0,
        )
    except Exception as e:
        return ModalityResult(
            modality="llm",
            raw_answer=None,
            canonical_bytes=None,
            error=f"LLM_ERROR:{type(e).__name__}",
            elapsed_ms=(time.monotonic() - t0) * 1000.0,
        )
    elapsed = (time.monotonic() - t0) * 1000.0
    raw_text = (raw or "").strip()
    canonical = _canonicalize_via_kernel(raw_text, pi_star_ref)
    if canonical is None:
        return ModalityResult(
            modality="llm",
            raw_answer=raw_text,
            canonical_bytes=None,
            error="PIS_REJECT",
            elapsed_ms=elapsed,
        )
    return ModalityResult(
        modality="llm",
        raw_answer=raw_text,
        canonical_bytes=canonical,
        error=None,
        elapsed_ms=elapsed,
    )


def run_witness(
    *,
    question: str,
    pi_star_ref: str,
    canonical_answer_bytes: bytes,
    cache_lookup: Callable[[], bytes | None] | None = None,
    chat_client: ChatClient | None = None,
    model_id: str = "",
    timeout_s: float = 10.0,
    progress=None,
) -> Witness:
    """Fan out kernel + cache + LLM in parallel; return a Witness.

    The kernel result is already known (passed in as
    ``canonical_answer_bytes``) — no work to do. Cache and LLM run in
    parallel threads. Total wall time ≈ max(cache, LLM) + small
    overhead, not sum.

    ``cache_lookup`` is a closure to keep this module independent of
    any storage concerns. Pass ``None`` (default) when no cache is
    available — modality reports ``ABSENT``.

    ``chat_client`` set to ``None`` skips the LLM modality entirely.

    ``timeout_s`` bounds the whole parallel stage. Modalities that
    don't finish in time report ``TIMEOUT``.
    """
    from arborist.qa.progress import disabled as _progress_disabled
    progress = progress or _progress_disabled()

    t_start = time.monotonic()
    progress.emit(
        "witness.start",
        parallel=2 + (1 if chat_client is not None else 0),
        timeout_s=timeout_s,
    )

    kernel_mod = _run_kernel_modality(canonical_answer_bytes, pi_star_ref)
    progress.emit(
        "witness.modality.kernel",
        ms=int(kernel_mod.elapsed_ms),
        ok=str(kernel_mod.ok).lower(),
    )

    modalities: dict[str, ModalityResult] = {"kernel": kernel_mod}

    # Submit cache + llm to the pool in parallel.
    with ThreadPoolExecutor(max_workers=2, thread_name_prefix="witness") as pool:
        fut_cache = pool.submit(
            _run_cache_modality, cache_lookup, pi_star_ref
        )
        fut_llm = pool.submit(
            _run_llm_modality, question, chat_client, model_id, pi_star_ref
        )

        try:
            cache_mod = fut_cache.result(timeout=timeout_s)
        except FutTimeoutError:
            cache_mod = ModalityResult(
                modality="cache",
                raw_answer=None,
                canonical_bytes=None,
                error="TIMEOUT",
                elapsed_ms=timeout_s * 1000.0,
            )
            fut_cache.cancel()
        modalities["cache"] = cache_mod
        progress.emit(
            "witness.modality.cache",
            ms=int(cache_mod.elapsed_ms),
            ok=str(cache_mod.ok).lower(),
            reason=(cache_mod.error or "ok"),
        )

        try:
            llm_mod = fut_llm.result(timeout=timeout_s)
        except FutTimeoutError:
            llm_mod = ModalityResult(
                modality="llm",
                raw_answer=None,
                canonical_bytes=None,
                error="TIMEOUT",
                elapsed_ms=timeout_s * 1000.0,
            )
            fut_llm.cancel()
        modalities["llm"] = llm_mod
        progress.emit(
            "witness.modality.llm",
            ms=int(llm_mod.elapsed_ms),
            ok=str(llm_mod.ok).lower(),
            reason=(llm_mod.error or "ok"),
        )

    label = _classify_agreement(modalities, canonical_answer_bytes)
    elapsed = (time.monotonic() - t_start) * 1000.0
    progress.emit("witness.done", label=label, ms=int(elapsed))

    return Witness(
        pi_star_ref=pi_star_ref,
        canonical_answer_bytes=canonical_answer_bytes,
        modalities=modalities,
        agreement_label=label,
        elapsed_ms=elapsed,
    )
