#!/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)


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)

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


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 = run_model(args.model, prompt_text, document_text, enable_thinking, sampling)
    finished = datetime.now(timezone.utc)

    reasoning, answer = split_reasoning(raw)

    # An empty answer is NOT "the checker found nothing". It means the model
    # produced only a reasoning trace, or nothing at all. Those are different
    # results and must never be recorded as a finding of silence — that is the
    # precise confusion trial 02 fell into. Flag it loudly and record it.
    degraded: str | None = None
    if not answer.strip():
        degraded = (
            "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."
        )
        print(f"\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()),
            "reasoning_present": reasoning is not None,
            "answer_words": len(answer.split()),
            "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()