#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import subprocess
import sys


def _parse_opencode_json_stream(raw_stdout: str) -> dict:
    text_parts: list[str] = []
    saw_finish = False
    malformed_lines = 0
    for line in raw_stdout.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            event = json.loads(line)
        except json.JSONDecodeError:
            malformed_lines += 1
            continue
        if event.get("type") == "text":
            part = event.get("part") or {}
            text = part.get("text")
            if isinstance(text, str):
                text_parts.append(text)
        if event.get("type") == "step_finish":
            saw_finish = True
    return {
        "text": "".join(text_parts).strip(),
        "saw_finish": saw_finish,
        "malformed_lines": malformed_lines,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Run opencode once via Python subprocess")
    parser.add_argument("--model", required=True)
    parser.add_argument("--agent", required=True)
    parser.add_argument("--cwd", required=False)
    parser.add_argument("--timeout-sec", type=int, default=600)
    parser.add_argument("--prompt", required=True)
    args = parser.parse_args()

    command = [
        "opencode",
        "run",
        "--pure",
        "--format",
        "json",
        "--model",
        args.model,
        "--agent",
        args.agent,
        args.prompt,
    ]

    try:
        completed = subprocess.run(
            command,
            cwd=args.cwd or None,
            capture_output=True,
            text=True,
            timeout=max(1, int(args.timeout_sec)),
            check=False,
        )
        parsed = _parse_opencode_json_stream(completed.stdout)
        completed_ok = bool(parsed["saw_finish"] or parsed["text"])
        payload = {
            "ok": completed.returncode == 0 and completed_ok,
            "returncode": completed.returncode,
            "stdout": parsed["text"],
            "stderr": completed.stderr,
            "timeout": False,
            "step_finish_seen": parsed["saw_finish"],
            "malformed_output_lines": parsed["malformed_lines"],
            "raw_stdout": completed.stdout,
        }
    except subprocess.TimeoutExpired as exc:
        raw_stdout = exc.stdout.decode() if isinstance(exc.stdout, bytes) else (exc.stdout or "")
        raw_stderr = exc.stderr.decode() if isinstance(exc.stderr, bytes) else (exc.stderr or "")
        parsed = _parse_opencode_json_stream(raw_stdout)
        timeout_ok = bool(parsed["saw_finish"] or (parsed["text"] and not raw_stderr.strip()))
        payload = {
            "ok": timeout_ok,
            "returncode": 124,
            "stdout": parsed["text"],
            "stderr": raw_stderr,
            "timeout": not timeout_ok,
            "step_finish_seen": parsed["saw_finish"],
            "malformed_output_lines": parsed["malformed_lines"],
            "raw_stdout": raw_stdout,
        }

    sys.stdout.write(json.dumps(payload))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())