#!/usr/bin/env python3
"""Reproduce OpenCode ACP cold and warm startup timings in a fresh Run9 box."""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from datetime import UTC, datetime
from pathlib import Path


DEFAULT_CONFIG = {
    "model": "chat/claude-standard",
    "enabled_providers": ["chat"],
    "provider": {
        "chat": {
            "npm": "@ai-sdk/amazon-bedrock",
            "options": {"region": "us-east-1"},
            "models": {
                "claude-standard": {
                    "id": (
                        "arn:aws:bedrock:us-east-1:841162690310:"
                        "application-inference-profile/2r34a06w99bz"
                    ),
                    "name": "Claude Standard",
                    "limit": {"context": 200000, "output": 16384},
                }
            },
        }
    },
}


def execute(
    command: list[str], *, capture: bool = False
) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(command, text=True, capture_output=capture, check=False)
    if result.returncode == 0:
        return result
    if capture:
        sys.stderr.write(result.stdout)
        sys.stderr.write(result.stderr)
    raise subprocess.CalledProcessError(result.returncode, command)


def config_content(path: str | None) -> str:
    if path:
        payload = json.loads(Path(path).read_text(encoding="utf-8"))
        return json.dumps(payload, separators=(",", ":"))
    payload = os.environ.get("OPENCODE_CONFIG_CONTENT")
    if payload:
        json.loads(payload)
        return payload
    return json.dumps(DEFAULT_CONFIG, separators=(",", ":"))


def parse_result(output: str) -> dict[str, object]:
    for line in reversed(output.splitlines()):
        try:
            value = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(value, dict) and "initialize_ms" in value:
            return value
    raise RuntimeError(f"ACP timing result was not found in output:\n{output}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Create a fresh Run9 box from a snap and run the OpenCode ACP "
            "initialize/session-new probe repeatedly. The first run is cold; "
            "later runs reuse the same box page cache."
        )
    )
    parser.add_argument(
        "--project",
        default=os.environ.get("RUN9_CLI_PROJECT"),
        required=os.environ.get("RUN9_CLI_PROJECT") is None,
        help="Run9 CLI project CID, or set RUN9_CLI_PROJECT",
    )
    parser.add_argument(
        "--snap",
        default=os.environ.get("RUN9_SOURCE_SNAP_ID"),
        required=os.environ.get("RUN9_SOURCE_SNAP_ID") is None,
        help="source snap ID, or set RUN9_SOURCE_SNAP_ID",
    )
    parser.add_argument("--box", help="box ID; defaults to a unique benchmark ID")
    parser.add_argument("--shape", default="1c2g")
    parser.add_argument("--binary", default="/usr/local/bin/opencode")
    parser.add_argument("--cwd", default="/workspace")
    parser.add_argument("--runs", type=int, default=2)
    parser.add_argument("--timeout", type=int, default=45)
    parser.add_argument(
        "--minimal",
        action="store_true",
        help="enable the experimental OPENCODE_ACP_MINIMAL startup path",
    )
    parser.add_argument(
        "--config-file",
        help=(
            "OpenCode JSON config. Defaults to a non-secret config with the same "
            "custom-provider shape as dev; OPENCODE_CONFIG_CONTENT is also accepted."
        ),
    )
    args = parser.parse_args()
    if args.runs < 1:
        parser.error("--runs must be at least 1")

    box = args.box or (
        "ib-acp-benchmark-"
        + datetime.now(UTC).strftime("%Y%m%d%H%M%S")
        + f"-{os.getpid()}"
    )
    box_cli = ["sys9", "run", "box"]
    probe = Path(__file__).with_name("smoke-opencode-acp.py").resolve()
    remote_probe = f"{args.cwd.rstrip('/')}/.acp-startup-smoke.py"

    print(f"Creating fresh box {box} from snap {args.snap}...", flush=True)
    execute(
        box_cli
        + [
            "create",
            box,
            "--project",
            args.project,
            "--snap",
            args.snap,
            "--shape",
            args.shape,
            "--description",
            "OpenCode ACP cold-start benchmark",
            "--label",
            "app=investibook",
            "--label",
            "purpose=acp-startup-benchmark",
        ]
    )
    execute(
        box_cli + ["cp", str(probe), f"{box}:{remote_probe}", "--project", args.project]
    )

    results: list[dict[str, object]] = []
    config = config_content(args.config_file)
    minimal_args = ["--acp-arg=--minimal"] if args.minimal else []
    for index in range(args.runs):
        label = "cold" if index == 0 else f"warm-{index}"
        print(f"\n=== {label} run ===", flush=True)
        started = time.monotonic()
        result = execute(
            box_cli
            + [
                "exec",
                box,
                "--project",
                args.project,
                "--deadline",
                f"{args.timeout + 15}s",
                "--workdir",
                args.cwd,
                "-e",
                "OPENCODE_ACP_PROFILE=1",
                "-e",
                "OPENCODE_DISABLE_DEFAULT_PLUGINS=1",
                "-e",
                f"OPENCODE_CONFIG_CONTENT={config}",
                "python3",
                remote_probe,
                "--binary",
                args.binary,
                "--cwd",
                args.cwd,
                "--timeout",
                str(args.timeout),
                *minimal_args,
            ],
            capture=True,
        )
        if result.stderr:
            sys.stderr.write(result.stderr)
        if result.stdout:
            sys.stdout.write(result.stdout)
        timing = parse_result(result.stdout)
        timing["run"] = label
        timing["box_exec_elapsed_ms"] = round((time.monotonic() - started) * 1000)
        results.append(timing)

    print("\n=== summary ===")
    print(json.dumps({"box": box, "snap": args.snap, "runs": results}, indent=2))
    print(
        f"\nBox retained for inspection. Remove it with:\n  sys9 run box rm {box} --project {args.project}"
    )


if __name__ == "__main__":
    main()