# /// script
# requires-python = ">=3.11"
# dependencies = ["chio-sdk"]
# ///
"""Run an actual repricing job and retain its independently checked trace/receipt association."""
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import socket
import subprocess
import time
import urllib.error
import urllib.request
import uuid

from chio.invariants import canonicalize_json, sha256_hex_utf8, verify_receipt_with_trusted_signers

ROOT = Path(__file__).resolve().parent
MODES = ["complete", "refuse", "fail", "uncertain"]


def request(path, token, body=None):
    query = urllib.request.Request(
        "http://127.0.0.1:4312" + path,
        headers={"Authorization": "Bearer " + token, "Content-Type": "application/json"},
        data=json.dumps(body).encode() if body is not None else None,
    )
    with urllib.request.urlopen(query, timeout=10) as response:
        return json.loads(response.read())


def verify(record, trusted_key):
    """Check signed identity and content bindings; OTLP alone is not signed evidence."""
    receipt = record["job"]["result"]["receipt"]
    checked = verify_receipt_with_trusted_signers(receipt, [trusted_key])
    if not checked["ok"]:
        raise ValueError("The job receipt is invalid or its signer is untrusted")
    span = record["otlp"]["resourceSpans"][0]["scopeSpans"][0]["spans"][0]
    attributes = {entry["key"]: entry["value"]["stringValue"] for entry in span["attributes"]}
    metadata = receipt["metadata"]
    expected = record["job"]["request_id"]
    if not (expected == record["job"]["job_id"] == attributes["chio.request.id"] == metadata["receipt_context"]["request_id"]):
        raise ValueError("Trace, job and receipt name different requests")
    if attributes["chio.receipt.id"] != receipt["id"] or record["job"]["result"]["receipt_id"] != receipt["id"]:
        raise ValueError("Trace names another receipt")
    if span["traceId"] != metadata["trace_id"] or span["spanId"] != metadata["span_id"]:
        raise ValueError("Trace identifiers differ from their signed receipt metadata")
    if attributes["chio.verdict"] != record["job"]["result"]["verdict"]:
        raise ValueError("Trace verdict differs from the retained outcome")
    if record["job"]["result"]["verdict"].lower() != receipt["decision"]["verdict"]:
        raise ValueError("Displayed outcome differs from the signed decision")
    if record["input"] != receipt["action"]["parameters"]:
        raise ValueError("Receipt describes a different repricing input")
    if sha256_hex_utf8(canonicalize_json(record["input"])) != receipt["action"]["parameter_hash"]:
        raise ValueError("Request parameter hash changed")
    if receipt["decision"]["verdict"] == "allow":
        if sha256_hex_utf8(canonicalize_json(record["job"]["result"]["output"])) != receipt["content_hash"]:
            raise ValueError("Price sheet output differs from its signed content hash")
    start, end = int(span["startTimeUnixNano"]), int(span["endTimeUnixNano"])
    if end < start:
        raise ValueError("Span ends before it starts")
    return {"ok": True, "signature": True, "request_binding": True, "trace_receipt_binding": True,
            "duration_ms": (end - start) / 1_000_000}


def run(binary, directory, mode, percent):
    state = directory / mode
    state.mkdir(parents=True, mode=0o700)
    environment = {**os.environ, "CHIO_APP_STATE": str(state.relative_to(ROOT))}
    with (state / "service.log").open("w") as log:
        process = subprocess.Popen([str(binary)], cwd=ROOT, env=environment, stdin=subprocess.DEVNULL, stdout=log, stderr=log)
        try:
            deadline = time.monotonic() + 60
            while time.monotonic() < deadline:
                if process.poll() is not None:
                    raise RuntimeError(f"Repricing service stopped; inspect {state}/service.log")
                if (state / "client-token.txt").exists() and (state / "kernel-public-key.txt").exists():
                    token = (state / "client-token.txt").read_text().strip()
                    try:
                        request("/budget", token)
                        break
                    except (OSError, ValueError):
                        pass
                time.sleep(.1)
            else:
                raise RuntimeError("Repricing service did not become ready")
            # Pin the public key from the kernel this operator owns, before requesting work.
            trusted_key = (state / "kernel-public-key.txt").read_text().strip()
            created = request("/jobs", token, {"percent": percent, "mode": mode})
            identity = created["job_id"]
            deadline = time.monotonic() + 60
            while time.monotonic() < deadline:
                job = request("/jobs/" + identity, token)
                trace = state / "jobs" / (identity + ".trace.json")
                if not job["active"] and job["result"] and trace.exists():
                    break
                time.sleep(.1)
            else:
                raise RuntimeError("The repricing job did not reach a retained outcome")
            record = {"mode": mode, "input": {"job_id": identity, "percent": percent, "mode": mode},
                      "job": job, "otlp": json.loads(trace.read_text()), "trusted_kernel_key": trusted_key,
                      "budget": request("/budget", token)}
            record["verification"] = verify(record, trusted_key)
            expected = {"complete": (True, True), "refuse": (False, False), "fail": (True, False), "uncertain": (True, True)}[mode]
            if (job["started"], job["published"]) != expected:
                raise ValueError("Price sheet effects do not match the selected path")
            (state / "trace-receipt.json").write_text(json.dumps(record, indent=2) + "\n")
            return record
        finally:
            if process.poll() is None:
                process.terminate()
            try:
                process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("percent", nargs="?", type=float, default=10.0)
    parser.add_argument("--mode", choices=MODES, default="complete")
    parser.add_argument("--compare", action="store_true", help="run all four outcomes with separate retained authority and data")
    args = parser.parse_args()
    if not 0 <= args.percent <= 100:
        parser.error("percent must be between 0 and 100")
    with socket.socket() as probe:
        try:
            probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            probe.bind(("127.0.0.1", 4312))
        except OSError:
            parser.error("Port 4312 is in use; stop the running repricing service first")
    subprocess.run(["cargo", "build", "--locked", "--bin", "repricing-service"], cwd=ROOT, check=True)
    target = Path(os.environ.get("CARGO_TARGET_DIR", ROOT / "target"))
    if not target.is_absolute():
        target = ROOT / target
    directory = ROOT / "state/trace-runs" / str(uuid.uuid4())
    os.umask(0o077)
    records = [run(target / "debug/repricing-service", directory, mode, args.percent) for mode in (MODES if args.compare else [args.mode])]
    (directory / "comparison.json").write_text(json.dumps(records, indent=2) + "\n")
    for record in records:
        print(f"{record['mode']}: {record['job']['result']['verdict']}; started={record['job']['started']}; published={record['job']['published']}; trace and receipt verified")
    print("Records:", directory.relative_to(ROOT) / "comparison.json")


if __name__ == "__main__":
    main()