"""Real-host state, claim and administrative exercises on fresh disposable state."""

import argparse
import json
import os
import shutil
import subprocess
import time
from pathlib import Path
from chio_process import ProcessClient, WorkerError
from chio_process.snapshot import JsonSnapshot

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--chio", type=Path, default=Path("chio"))
parser.add_argument("--state", type=Path, required=True)
args = parser.parse_args()
os.umask(0o077)
root = args.state.resolve()
root.mkdir(mode=0o700)
chio = str(Path(shutil.which(str(args.chio)) or str(args.chio)).resolve(strict=True))


def command(*parts):
    return json.loads(
        subprocess.check_output(
            [chio, "--json", "process", *map(str, parts)], text=True
        )
    )


def write(name, value):
    path = root / name
    path.write_text(json.dumps(value))
    return path


(root / "policy.yaml").write_text(
    """kernel:
  max_capability_ttl: 3600
  delegation_depth_limit: 1
  durable_admission_mode: all
capabilities:
  default:
    tools:
      - server: chio-ipc
        tool: '*'
        operations: [invoke, delegate]
        ttl: 3600
"""
)
routes = [
    "send_jobs",
    "receive_jobs",
    "ack_jobs",
    "claim_jobs",
    "complete_jobs",
    "renew_jobs",
]
config = write(
    "config.json",
    {
        "schema": "chio.process.host.v1",
        "policy": str(root / "policy.yaml"),
        "mailboxes": [
            {
                "id": "jobs",
                "renewable_leases": True,
                "limits": {
                    "max_pending_messages": 1,
                    "max_pending_bytes": 65536,
                    "max_message_bytes": 65536,
                    "max_messages": 2,
                },
            }
        ],
        "limits": {"max_processes": 3, "max_depth": 1, "max_calls": 32},
        "children": [
            {
                "id": role,
                "parent": "root",
                "budget_share_bps": 5000,
                "tools": [
                    {"server_id": "chio-ipc", "tool_name": route} for route in routes
                ],
            }
            for role in ["a", "b"]
        ],
    },
)
host = root / "host"
initialized = command("init", "--state", host, "--config", config)
connections = {}
socket = root / "sock"
for role in ["a", "b"]:
    command(
        "credential",
        "--state",
        host,
        "--process",
        role,
        "--socket",
        socket,
        "--out",
        root / f"{role}.json",
    )
    connections[role] = json.loads((root / f"{role}.json").read_text())


def start():
    log = (root / "serve.log").open("ab")
    process = subprocess.Popen(
        [chio, "process", "serve", "--state", str(host), "--socket", str(socket)],
        stdout=log,
        stderr=log,
    )
    log.close()
    for _ in range(100):
        if socket.exists():
            return process
        if process.poll() is not None:
            raise RuntimeError("Host failed; inspect serve.log")
        time.sleep(0.05)
    process.terminate()
    raise TimeoutError("Host did not bind")


def client(role):
    c = connections[role]
    return ProcessClient(c["socket_path"], c["credential"])


def value(response):
    assert response["verdict"] == "allow", response["reason"]
    assert response["terminal_state"]["state"] == "completed"
    assert response["output"]["kind"] == "value"
    return response["output"]["value"]


process = start()
evidence = {}
try:
    a, b = client("a"), client("b")
    original = a.inspect()["checkpoint"]
    committed = a.checkpoint(
        original["revision"], {"report.v1": {"completed": ["one"]}}
    )
    try:
        a.checkpoint(original["revision"], {"report.v1": {"completed": ["two"]}})
    except WorkerError as error:
        assert error.code == "checkpoint_conflict"
    else:
        raise AssertionError("Stale CAS accepted")
    assert a.inspect()["checkpoint"] == committed
    snapshot = JsonSnapshot(a)
    document = {"report": ["retained-" * 1000] * 20}
    reference = snapshot.write(document)
    assert snapshot.read(reference) == document
    blob = a.put_blob(b"original report")
    assert a.read_blob(blob["sha256"]) == b"original report"
    try:
        b.read_blob(blob["sha256"])
    except WorkerError as error:
        assert error.code == "blob_missing"
    else:
        raise AssertionError("Sibling read another process blob")
    evidence["checkpoint_conflict_and_blob_isolation"] = True
    empty = a.invoke(
        "poll-empty", "chio-ipc", "receive_jobs", {"after_sequence": "0", "limit": 1}
    )
    assert value(empty)["messages"] == []
    sent = a.invoke(
        "send-one",
        "chio-ipc",
        "send_jobs",
        {"message_key": "one", "payload": {"items": [2, 3, 5]}},
    )
    assert value(sent)["status"] == "sent"
    full = value(
        a.invoke(
            "send-two-full",
            "chio-ipc",
            "send_jobs",
            {"message_key": "two", "payload": {}},
        )
    )
    assert full["status"] == "full"
    assert (
        a.invoke(
            "poll-empty",
            "chio-ipc",
            "receive_jobs",
            {"after_sequence": "0", "limit": 1},
        )["receipt_json"]
        == empty["receipt_json"]
    )
    messages = value(
        a.invoke(
            "poll-next", "chio-ipc", "receive_jobs", {"after_sequence": "0", "limit": 1}
        )
    )["messages"]
    assert messages[0]["sender"] == "a"
    claimed = value(
        a.invoke("claim-a", "chio-ipc", "claim_jobs", {"limit": 1, "lease_ms": 1000})
    )["messages"][0]
    renewal_args = {
        "sequence": claimed["sequence"],
        "claim": claimed["claim"],
        "lease_ms": 1000,
    }
    renewal = a.invoke("renew-a-0", "chio-ipc", "renew_jobs", renewal_args)
    assert (
        a.invoke("renew-a-0", "chio-ipc", "renew_jobs", renewal_args)["receipt_json"]
        == renewal["receipt_json"]
    )
    time.sleep(1.2)
    replacement = value(
        b.invoke("claim-b", "chio-ipc", "claim_jobs", {"limit": 1, "lease_ms": 30000})
    )["messages"][0]
    stale = a.invoke(
        "complete-stale",
        "chio-ipc",
        "complete_jobs",
        {"sequence": claimed["sequence"], "claim": claimed["claim"]},
    )
    # Native invalid completion may surface as a signed denial or completed tool error.
    assert (
        stale["verdict"] != "allow" or stale["output"]["value"].get("isError") is True
    )
    done = value(
        b.invoke(
            "complete-b",
            "chio-ipc",
            "complete_jobs",
            {"sequence": replacement["sequence"], "claim": replacement["claim"]},
        )
    )
    evidence["stale_completion"] = {
        "verdict": stale["verdict"],
        "terminal_state": stale["terminal_state"],
    }
    evidence["replacement_completion"] = done
    second = value(
        a.invoke(
            "send-two", "chio-ipc", "send_jobs", {"message_key": "two", "payload": {}}
        )
    )
    assert second["status"] == "sent"
    value(
        a.invoke(
            "ack-two", "chio-ipc", "ack_jobs", {"through_sequence": second["sequence"]}
        )
    )
    exhausted = value(
        a.invoke(
            "send-three",
            "chio-ipc",
            "send_jobs",
            {"message_key": "three", "payload": {}},
        )
    )
    assert exhausted["status"] == "exhausted"
    cursor = value(
        a.invoke(
            "expired-cursor",
            "chio-ipc",
            "receive_jobs",
            {"after_sequence": "0", "limit": 1},
        )
    )
    assert cursor["status"] == "cursor_expired"
    evidence.update(
        empty_poll_replayed=True,
        full_then_sent=True,
        lifetime_exhaustion=True,
        expired_cursor=True,
        inspect=a.inspect(),
        original_send_receipt=sent["receipt_json"],
    )
finally:
    process.terminate()
    process.wait(timeout=15)
# Admin commands deliberately run only after host shutdown.
command("revoke", "--state", host, "--process", "a")
command("cancel", "--state", host, "--process", "b")
old_credential = connections["a"]["credential"]
socket = root / "socket-next"
command(
    "credential",
    "--state",
    host,
    "--process",
    "a",
    "--socket",
    socket,
    "--out",
    root / "a-next.json",
)
connections["a"] = json.loads((root / "a-next.json").read_text())
process = start()
try:
    try:
        ProcessClient(str(socket), old_credential).inspect()
    except WorkerError as error:
        assert error.code == "unauthenticated"
    else:
        raise AssertionError("Revoked credential authenticated")
    a = client("a")
    replay = a.invoke(
        "send-one",
        "chio-ipc",
        "send_jobs",
        {"message_key": "one", "payload": {"items": [2, 3, 5]}},
    )
    assert replay["receipt_json"] == evidence["original_send_receipt"]
    from chio_process.invocation import invoke_recorded

    bound = invoke_recorded(
        chio,
        connections["a"],
        initialized["kernel_key"],
        {
            "operation_key": "verified-empty-poll",
            "server_id": "chio-ipc",
            "tool_name": "receive_jobs",
            "arguments": {"after_sequence": "2", "limit": 1},
            "known_outcome_only": True,
        },
        root / "bound-response",
    )
    verification = json.loads((root / "bound-response/verification.json").read_text())
    assert verification["response_bound"] is True
    evidence.update(
        revoked_credential_refused=True,
        rotation_preserved_receipt=True,
        response_bound=True,
    )
finally:
    process.terminate()
    process.wait(timeout=15)
evidence["offline_cancelled_state"] = command(
    "state", "--state", host, "--process", "b"
)
command("export", "--state", host)
moved = root / "moved"
shutil.copytree(host, moved)
command("import", "--state", moved)
evidence["relocated_state"] = command("state", "--state", moved, "--process", "a")
write("evidence.json", evidence)
print(
    json.dumps(
        {
            key: val
            for key, val in evidence.items()
            if key
            not in [
                "inspect",
                "original_send_receipt",
                "offline_cancelled_state",
                "relocated_state",
            ]
        },
        indent=2,
    )
)