#!/usr/bin/env python3
"""
Fool trial harness — reproducible runs of a differently-formed checker.

WHY THIS EXISTS
    Trials 01 and 02 were run ad hoc. What survives of them is the model ID, the
    MLX version, the hardware, `enable_thinking`, and the prompt *in paraphrase*.
    Never recorded: the verbatim prompt, temperature, top_p, max_tokens, seed.
    There was no script. Trial 03 therefore could not have been run under trial
    02's conditions, because those conditions were never written down.

    The same failure destroyed the v1 Chamber's GPT-side protocol: it lived as
    configuration inside a hosted product, was updated in place, and is gone.
    The Claude-side prompt survived because it was a file in a repository.

    A protocol held as ephemeral runtime state is not a protocol. This harness
    makes every run a file.

WHAT IT GUARANTEES
    - The prompt is a versioned file, hashed into the run record.
    - Every sampling parameter is recorded, including the defaults.
    - Reasoning trace and answer are separated but neither is suppressed
      (trial 02 established that disabling thinking makes the model mute,
      not terse — so `enable_thinking` defaults ON and turning it off is loud).
    - The raw, unparsed output is always kept.
    - Failure to load the model is an error, never an empty result. A trial that
      silently returns nothing is indistinguishable from a checker that found
      nothing, which is the one confusion this instrument cannot afford.

USAGE
    ./run_trial.py --trial 03 \
        --prompt prompts/trial-03-self-exemption.txt \
        --input  /path/to/document.md \
        --note   "self-exemption test: does the checker exempt its own justification?"

    Runs on the machine holding the model (CapableHands M4). Writes to runs/.
"""

from __future__ import annotations

import argparse
import getpass
import hashlib
import json
import platform
import re
import socket
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
RUNS_DIR = HERE / "runs"

DEFAULT_MODEL = "mlx-community/Qwen3.6-35B-A3B-8bit"

# Defaults are recorded in every run record even when unchanged, so that a later
# reader never has to ask what the harness "would have" used.
DEFAULT_SAMPLING = {
    "temperature": 0.7,
    "top_p": 0.95,
    "max_tokens": 4096,
    "seed": None,  # None = library default (non-deterministic); record it as such
}


def sha256(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def read_text(path: Path) -> str:
    if not path.is_file():
        sys.exit(f"FATAL: not a file: {path}")
    return path.read_text(encoding="utf-8")


def _package_version(dist_name: str, module_name: str, attr_module: str | None = None) -> str:
    """
    Resolve a package version without depending on a top-level __version__.

    `mlx` has no `mlx.__version__` — only `mlx.core.__version__` — so the naive
    getattr fallback silently records "unknown" for an installed, versioned
    package. That is the exact confusion this harness exists to prevent: a probe
    that could not look, reporting a value that reads like a result. Installed
    metadata is asked first; the module attribute is the fallback; and a genuine
    absence is reported as absence, distinguishably.
    """
    try:
        import importlib.metadata as md

        return md.version(dist_name)
    except Exception:
        pass
    try:
        import importlib

        mod = importlib.import_module(attr_module or module_name)
        v = getattr(mod, "__version__", None)
        if v:
            return str(v)
        return "INSTALLED, VERSION UNDETERMINED"
    except Exception as exc:  # pragma: no cover - environment probe
        return f"UNAVAILABLE ({exc.__class__.__name__})"


def environment() -> dict:
    """Capture enough of the machine to make a later re-run comparable."""
    return {
        "host": socket.gethostname(),
        "user": getpass.getuser(),
        "platform": platform.platform(),
        "machine": platform.machine(),
        "python": sys.version.split()[0],
        "mlx_version": _package_version("mlx", "mlx", attr_module="mlx.core"),
        "mlx_lm_version": _package_version("mlx-lm", "mlx_lm"),
    }


def harness_sha256() -> str:
    """
    Hash this file into every run record.

    `git_revision()` returns None whenever the harness is executed outside its
    repository — which is the normal case, because it must run on the machine
    holding the model. A run record whose harness field is null cannot be traced
    to the code that produced it. The prompt and input are already hashed; the
    instrument itself was not.
    """
    try:
        return sha256(Path(__file__).resolve().read_text(encoding="utf-8"))
    except Exception as exc:  # pragma: no cover
        return f"UNAVAILABLE ({exc.__class__.__name__})"


def git_revision() -> str | None:
    """Record which revision of the harness and prompts produced this run."""
    try:
        out = subprocess.run(
            ["git", "-C", str(HERE), "rev-parse", "--short", "HEAD"],
            capture_output=True, text=True, timeout=10,
        )
        return out.stdout.strip() or None
    except Exception:
        return None


THINK_RE = re.compile(r"<think>(.*?)</think>", re.DOTALL | re.IGNORECASE)

# Qwen3.6 does not always tag its scratchpad. In trial 03 it opened with the bare
# line "Here's a thinking process:" and never emitted a <think> tag, so the tag
# regex reported reasoning_present=false and the whole deliberation was recorded
# as the answer. These are openings of *deliberation about the task*, which no
# answer to this prompt begins with — the prompt forbids summarising the document
# and asks for named assumptions.
UNTAGGED_SCRATCHPAD_RE = re.compile(
    r"^\s*(?:"
    # First-person statements of intent about the task.
    r"(?:here(?:'|’)s|here is|let(?:'|’)s|i(?:'|’)ll|i will|i need to|i should|"
    r"first,?\s+i)\b"
    # Interjections, which must actually be interjections. A bare `okay` matched
    # "Okay is not a word this document uses, but its approach…" — a sentence that
    # belongs in an answer. The punctuation is what distinguishes the two.
    r"|(?:okay|ok|alright|right|so)\s*[,:]"
    r")[^\n]{0,80}"
    r"(?:thinking process|thought process|think through|reasoning|analyz|approach|"
    r"plan\b|break (?:this|it) down|work through|go section by section)",
    re.IGNORECASE,
)


def split_reasoning(raw: str) -> tuple[str | None, str]:
    """
    Separate the reasoning trace from the answer WITHOUT discarding either.

    Returns (reasoning_or_None, answer). If no trace is present the whole output
    is the answer and reasoning is None — recorded as such rather than inferred.
    """
    matches = THINK_RE.findall(raw)
    if not matches:
        return None, raw.strip()
    reasoning = "\n\n---\n\n".join(m.strip() for m in matches)
    answer = THINK_RE.sub("", raw).strip()
    return reasoning, answer


def run_model(
    model_id: str,
    prompt: str,
    document: str,
    enable_thinking: bool,
    sampling: dict,
) -> str:
    """Load the model and generate. Any failure is fatal and loud."""
    try:
        from mlx_lm import generate, load
        from mlx_lm.sample_utils import make_sampler
    except ImportError as exc:
        sys.exit(
            f"FATAL: mlx_lm unavailable ({exc}).\n"
            "This harness must run on the machine holding the model."
        )

    model, tokenizer = load(model_id)

    messages = [
        {"role": "system", "content": prompt},
        {"role": "user", "content": document},
    ]

    # Qwen-family templates accept enable_thinking; older templates do not.
    # Try the explicit form, and record if we had to fall back.
    try:
        text = tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, enable_thinking=enable_thinking
        )
    except TypeError:
        if not enable_thinking:
            sys.exit(
                "FATAL: enable_thinking=False requested but this tokenizer's chat "
                "template does not support it. Refusing to run — a silent fallback "
                "to thinking-on would misattribute the result."
            )
        text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)

    sampler_kwargs = {"temp": sampling["temperature"], "top_p": sampling["top_p"]}
    sampler = make_sampler(**sampler_kwargs)

    raw = generate(
        model,
        tokenizer,
        prompt=text,
        max_tokens=sampling["max_tokens"],
        sampler=sampler,
        verbose=False,
    )

    # Re-encoding the decoded text is an ESTIMATE, not the true generated count —
    # encode(decode(x)) is not guaranteed to round-trip. It is reported as an
    # estimate and only used to detect the token ceiling, where being a few tokens
    # out cannot change the verdict.
    try:
        generated_tokens = len(tokenizer.encode(raw))
    except Exception:
        generated_tokens = None

    return raw, generated_tokens


def main() -> None:
    ap = argparse.ArgumentParser(description="Run one Fool trial, reproducibly.")
    ap.add_argument("--trial", required=True, help="Trial number, e.g. 03")
    ap.add_argument("--prompt", required=True, type=Path, help="Prompt file (versioned)")
    ap.add_argument("--input", required=True, type=Path, help="Document under test")
    ap.add_argument("--note", default="", help="One line: what this trial tests")
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--temperature", type=float, default=DEFAULT_SAMPLING["temperature"])
    ap.add_argument("--top-p", type=float, default=DEFAULT_SAMPLING["top_p"])
    ap.add_argument("--max-tokens", type=int, default=DEFAULT_SAMPLING["max_tokens"])
    ap.add_argument("--seed", type=int, default=DEFAULT_SAMPLING["seed"])
    ap.add_argument(
        "--no-thinking",
        action="store_true",
        help=(
            "Disable the reasoning mode. Trial 02 established this makes the model "
            "MUTE, not terse. Only use to deliberately re-measure that."
        ),
    )
    args = ap.parse_args()

    enable_thinking = not args.no_thinking
    if not enable_thinking:
        print(
            "WARNING: enable_thinking=False. Trial 02 found this produces silence, "
            "not brevity. A `nothing found` result from this run is UNINTERPRETABLE "
            "as restraint.",
            file=sys.stderr,
        )

    prompt_text = read_text(args.prompt.resolve())
    document_text = read_text(args.input.resolve())

    sampling = {
        "temperature": args.temperature,
        "top_p": args.top_p,
        "max_tokens": args.max_tokens,
        "seed": args.seed,
    }

    if args.seed is not None:
        try:
            import mlx.core as mx

            mx.random.seed(args.seed)
        except Exception as exc:
            sys.exit(f"FATAL: seed requested but could not be set ({exc}).")

    started = datetime.now(timezone.utc)
    raw, generated_tokens = run_model(
        args.model, prompt_text, document_text, enable_thinking, sampling
    )
    finished = datetime.now(timezone.utc)

    reasoning, answer = split_reasoning(raw)

    # A result is degraded whenever what was recorded as `answer` is not an answer.
    #
    # The original guard tested only for emptiness, which is a property of the
    # STRING while the field claims a property of the RESULT. Trial 03 walked
    # straight through it: 2,944 words of untagged deliberation, cut off at the
    # token ceiling before the answer began, recorded as `degraded: null`. Trial 02
    # mistook silence for restraint; that guard would have let trial 03 mistake
    # deliberation for a finding. Every condition below must therefore be reported,
    # not just the first.
    problems: list[str] = []

    if not answer.strip():
        problems.append(
            "EMPTY ANSWER: the model produced no text outside its reasoning trace. "
            "This is a harness/generation failure, NOT a finding of 'nothing found'. "
            "Do not grade it as restraint."
        )

    ceiling = sampling["max_tokens"]
    hit_ceiling = generated_tokens is not None and generated_tokens >= ceiling - 2
    untagged_scratchpad = reasoning is None and bool(UNTAGGED_SCRATCHPAD_RE.match(answer))

    if hit_ceiling:
        problems.append(
            f"TOKEN CEILING: generation stopped at the max_tokens limit "
            f"(~{generated_tokens} of {ceiling}). The output is CUT OFF, not "
            f"complete. Anything absent from it may simply never have been reached."
        )

    # Trials 03 and 04 are the two halves of this distinction, and the first
    # version of the guard collapsed them. Trial 03: untagged deliberation that ran
    # into the ceiling, so no answer ever existed and the run was VOID. Trial 04:
    # untagged deliberation that COMPLETED, so the answer follows the scratchpad in
    # the same file and the run is perfectly gradeable once it is extracted.
    #
    # Reporting both as "do not grade this" would have thrown away six good runs.
    # Reporting both as fine would have graded trial 03's silence as restraint.
    if untagged_scratchpad and hit_ceiling:
        problems.append(
            "VOID — DELIBERATION, THEN TRUNCATION: the output opens as deliberation "
            "about the task, no <think> tag was emitted, and generation hit the "
            "ceiling before any answer began. NOTHING here is a finding, and the "
            "absence of findings is NOT restraint. Re-run with a higher ceiling."
        )
    elif untagged_scratchpad:
        problems.append(
            "ANSWER EMBEDDED: the output opens as deliberation and no <think> tag "
            "was emitted, so the scratchpad was recorded as the ANSWER — but "
            "generation COMPLETED, so the real answer follows it in the same file. "
            "This run is NOT void. Extract the answer before grading, and do not "
            "count the deliberation as findings."
        )

    # `think_tag_found: false` says a regex found no tag. It has never said the
    # model did not deliberate, and trial 04 is the proof: 2,944 words of it.

    degraded: str | None = "\n".join(problems) if problems else None
    if degraded:
        print(f"\n*** DEGRADED RUN ***\n{degraded}\n", file=sys.stderr)

    stamp = started.strftime("%Y%m%dT%H%M%SZ")
    slug = f"trial-{args.trial}-{stamp}"
    RUNS_DIR.mkdir(parents=True, exist_ok=True)

    record = {
        "trial": args.trial,
        "note": args.note,
        "started_utc": started.isoformat(),
        "finished_utc": finished.isoformat(),
        "duration_s": round((finished - started).total_seconds(), 1),
        "model": args.model,
        "enable_thinking": enable_thinking,
        "sampling": sampling,
        "prompt": {
            "path": str(args.prompt),
            "sha256": sha256(prompt_text),
            "words": len(prompt_text.split()),
        },
        "input": {
            "path": str(args.input),
            "sha256": sha256(document_text),
            "words": len(document_text.split()),
        },
        "output": {
            "raw_words": len(raw.split()),
            # Named for what it actually tests. The old key was `reasoning_present`,
            # which read as a claim about the model and was in fact a claim about a
            # regex: trial 03 deliberated for 2,944 words and this field said false.
            "think_tag_found": reasoning is not None,
            "answer_words": len(answer.split()),
            "generated_tokens_est": generated_tokens,
            "hit_token_ceiling": (
                None if generated_tokens is None
                else generated_tokens >= sampling["max_tokens"] - 2
            ),
            "degraded": degraded,
        },
        "environment": environment(),
        "harness_git_rev": git_revision(),
        "harness_sha256": harness_sha256(),
    }

    (RUNS_DIR / f"{slug}.json").write_text(
        json.dumps(record, indent=2) + "\n", encoding="utf-8"
    )
    # The raw output is kept verbatim and unparsed. The split below is a
    # convenience view; the raw file is the record of what was actually produced.
    (RUNS_DIR / f"{slug}.raw.txt").write_text(raw, encoding="utf-8")
    (RUNS_DIR / f"{slug}.answer.md").write_text(answer + "\n", encoding="utf-8")
    if reasoning is not None:
        (RUNS_DIR / f"{slug}.reasoning.md").write_text(reasoning + "\n", encoding="utf-8")

    print(f"\n{'=' * 60}")
    print(f"trial {args.trial} · {record['duration_s']}s · {record['output']['answer_words']} words")
    print(f"reasoning trace: {'captured' if reasoning else 'ABSENT — check template'}")
    print(f"records → {RUNS_DIR / slug}.*")
    print(f"{'=' * 60}\n")
    print(answer)


if __name__ == "__main__":
    main()