"""J-Space cognition layer — inference-time control protocol from the
J-Space Cognition Suite (github.com/Tiger3807861189/J-Space-Cognition-Suite-V3.6),
adapted to koi-harness's router/planner/DAG shape.

Principle: "dense on the inside, decodable on demand, clean on the outside."
Machinery loads only when task complexity justifies it — fast mode adds zero overhead.
"""
from __future__ import annotations

import re
import time
from pathlib import Path

# ── mode gate (J-Space pass criteria) ────────────────────────────
# fast: one step, or a step you can check in one glance → no machinery
# full: two to four steps, one deliverable, verifiable in one reading → hub + verify
# loop: multiple stages/files/turns → hub + ledger + verify
MODE_GATE_RULES = (
    'Also classify cognition depth as "cog": '
    '"fast" (one step, or checkable in a single glance — if you cannot verify it in '
    'one glance it is NOT fast), '
    '"full" (two to four dependent steps, one deliverable, verifiable in one reading), '
    '"loop" (multiple stages, multiple files, or work spanning many turns).'
)

# ── inner register rules for LLM workers (dense track) ───────────
DENSE_TRACK = (
    " Keep an inner working register while you work: before acting, name the one or two "
    "ideas currently staged and the single fact that makes each matter. Track epistemic "
    "state per item — verified (only with its scope stated), uncertain, or refuted — and "
    "never re-derive something you already established; reuse it. Derive intermediate "
    "facts BEFORE the conclusions that depend on them. Your final answer is the outer "
    "register: clean, complete language only — no shorthand, no symbols, no half-compressed "
    "notes leaking through."
)

# ── broadcast hub instruction for the planner ────────────────────
HUB_RULES = (
    ' Additionally emit "hub": a list of at most 5 short strings — the shared names, values, '
    "and constraints that multiple steps depend on, each stated with the one fact that makes "
    "it matter. Derive each ONCE here; they are broadcast to every step so no step re-derives them."
)

# ── verification gate (adapted J-Space invariants) ───────────────
VERIFIER_SYSTEM = (
    "You are a verification gate. Re-read the task goal line by line and audit the candidate "
    "answer against these invariants: (1) every part of the goal is actually addressed — a task "
    "is never finished without this re-read; (2) no unresolved placeholder, shorthand, or "
    "half-compressed fragment leaked into the answer; (3) no internal contradiction between parts "
    "of the answer; (4) claims that came from different steps agree with each other. "
    "You may ONLY flag issues fixable from the provided step outputs. The step outputs are the "
    "ground truth of what happened — never demand external artifacts (logs, screenshots, URLs, "
    "re-verification) and never question whether the work was really performed. "
    'Reply ONLY with JSON: {"ok": true} or {"ok": false, "issues": ["specific, actionable issue", ...]}.'
)

REPAIR_SYSTEM = (
    "Repair the candidate answer so the listed issues are resolved, using the provided step "
    "outputs and context — do not invent new facts. If an issue cannot be fixed from the "
    "evidence, leave that part of the answer as it was. NEVER replace the answer with a refusal, "
    "apology, or demand for more evidence — always return the best complete answer. Return the "
    "corrected answer only, in clean complete language."
)


class Ledger:
    """Durable state register: short labelled lines, re-read at every seam.
    Fields per J-Space: goal, core, verified, open, next."""

    def __init__(self, root: Path, goal: str):
        self.path = root / "workspaces" / "_ledger" / f"run-{int(time.time())}.md"
        self.goal = goal
        self.core: list[str] = []
        self.verified: list[str] = []
        self.open: list[str] = []
        self.next: str = ""
        self._flush()

    def seam(self, step_id: int, agent: str, result: str, remaining: list[int]) -> None:
        """Called when a step settles — the J-Space 'marker ends with settlement' rule."""
        head = re.sub(r"\s+", " ", result).strip()[:160]
        self.verified.append(f"step {step_id} ({agent}): {head}")
        self.open = [f"step {i}" for i in remaining]
        self.next = f"steps {remaining}" if remaining else "synthesize + verify"
        self._flush()

    def set_core(self, hub: list[str]) -> None:
        self.core = hub[:5]
        self._flush()

    def render(self) -> str:
        lines = [f"goal: {self.goal}"]
        lines += [f"core: {c}" for c in self.core]
        lines += [f"verified: {v}" for v in self.verified]
        lines += [f"open: {o}" for o in self.open]
        lines.append(f"next: {self.next}")
        return "\n".join(lines)

    def _flush(self) -> None:
        try:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            self.path.write_text(self.render() + "\n", encoding="utf-8")
        except OSError:
            pass  # ledger is an aid, never a failure source


async def verify_and_repair(llm, task: str, answer: str, evidence: str, max_repairs: int = 1) -> str:
    """Post-synthesis verification gate; one repair pass on failure."""
    for _ in range(max_repairs):
        try:
            verdict = await llm.json(
                VERIFIER_SYSTEM,
                f"Task goal:\n{task}\n\nCandidate answer:\n{answer}",
                max_tokens=512,
                temperature=0.0,
            )
        except Exception:
            return answer
        if verdict.get("ok", True):
            return answer
        issues = "\n".join(f"- {i}" for i in verdict.get("issues", []))
        answer = await llm.text(
            REPAIR_SYSTEM,
            f"Task goal:\n{task}\n\nIssues:\n{issues}\n\nCandidate answer:\n{answer}\n\n"
            f"Step outputs / context:\n{evidence}",
            max_tokens=4096,
        )
    return answer