#!/usr/bin/env python3
"""Run the native Swift host and compare its EOS-bound output with a BF16 receipt."""

from __future__ import annotations

import argparse
import json
import subprocess
import time
from pathlib import Path


EOS_TOKEN_ID = 2


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--executable", type=Path, required=True)
    parser.add_argument("--model-dir", type=Path, required=True)
    parser.add_argument("--pixel-values", type=Path, required=True)
    parser.add_argument("--expected-receipt", type=Path, required=True)
    parser.add_argument("--expected-row", type=int, default=0)
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    expected_receipt = json.loads(args.expected_receipt.read_text())
    if "rows" in expected_receipt:
        row = expected_receipt["rows"][args.expected_row]
        expected_tokens = row["reference_tokens"]
        expected_text = row["reference_text"]
        source = row.get("image")
    else:
        expected_tokens = expected_receipt.get("native_tokens", expected_receipt["coreml_tokens"])
        expected_tokens = expected_tokens[: expected_tokens.index(EOS_TOKEN_ID) + 1]
        expected_text = None
        source = "synthetic-canary"

    command = [
        str(args.executable.expanduser().resolve()),
        "--model-dir",
        str(args.model_dir.expanduser().resolve()),
        "--pixel-values-fp32",
        str(args.pixel_values.expanduser().resolve()),
        "--max-tokens",
        str(args.max_new_tokens),
    ]
    started = time.perf_counter()
    completed = subprocess.run(command, check=False, capture_output=True, text=True)
    seconds = time.perf_counter() - started
    stdout_lines = completed.stdout.splitlines()
    if not stdout_lines:
        raise RuntimeError(f"Swift runtime produced no stdout: {completed.stderr}")
    tokens = [int(value) for value in stdout_lines[0].split()]
    text = "\n".join(stdout_lines[1:])
    receipt = {
        "source": source,
        "command": [
            args.executable.name,
            "--model-dir",
            ".",
            "--pixel-values-fp32",
            args.pixel_values.name,
            "--max-tokens",
            str(args.max_new_tokens),
        ],
        "exit_code": completed.returncode,
        "seconds": seconds,
        "tokens": tokens,
        "expected_tokens": expected_tokens,
        "token_exact": tokens == expected_tokens,
        "stop_reason": "eos" if tokens and tokens[-1] == EOS_TOKEN_ID else "max_new_tokens",
        "text": text,
        "expected_text": expected_text,
        "text_exact": expected_text is None or text == expected_text,
        "stderr": completed.stderr,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(receipt, ensure_ascii=False, indent=2) + "\n")
    print(json.dumps(receipt, ensure_ascii=False, indent=2))
    if (
        completed.returncode != 0
        or not receipt["token_exact"]
        or not receipt["text_exact"]
        or receipt["stop_reason"] != "eos"
    ):
        raise SystemExit(1)


if __name__ == "__main__":
    main()