#!/usr/bin/env python3
"""Run a command-shaped Agent with the agent-continuity lifecycle around it."""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path


CLI_PATH = Path(__file__).with_name("agent-continuity.py")


def read_prompt(args: argparse.Namespace) -> str:
    supplied = [args.text is not None, args.input_file is not None, args.input_value is not None]
    if sum(supplied) != 1:
        raise ValueError("provide exactly one of --text, --input-file, or --input")
    if args.text is not None:
        return args.text
    if args.input_file is not None:
        return Path(args.input_file).expanduser().read_text(encoding="utf-8")
    if args.input_value == "-":
        return sys.stdin.read()
    return args.input_value


def run_cli(home: str | None, *arguments: str) -> dict:
    command = [sys.executable, str(CLI_PATH)]
    if home:
        command.extend(["--home", home])
    command.extend(arguments)
    result = subprocess.run(command, capture_output=True, text=True, encoding="utf-8")
    if result.returncode != 0:
        raise RuntimeError(result.stderr.strip() or "agent-continuity command failed")
    return json.loads(result.stdout) if result.stdout.strip() else {}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="wrap a command-shaped Agent with agent-continuity")
    parser.add_argument("--home")
    parser.add_argument("--session-id", required=True)
    parser.add_argument("--agent", required=True)
    parser.add_argument("--provider")
    parser.add_argument("--model")
    parser.add_argument("--lease-ttl-seconds", type=int)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--text")
    group.add_argument("--input-file")
    group.add_argument("--input", dest="input_value")
    parser.add_argument("command", nargs=argparse.REMAINDER, help="command after --")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    command = list(args.command)
    if command and command[0] == "--":
        command.pop(0)
    if not command:
        parser.error("provide an Agent command after --")

    prompt = read_prompt(args)
    temp_dir = Path(tempfile.mkdtemp(prefix="agent-continuity-run-"))
    prompt_path = temp_dir / "user-message.txt"
    response_path = temp_dir / "assistant-response.txt"
    stderr_path = temp_dir / "agent-stderr.txt"
    prompt_path.write_text(prompt, encoding="utf-8")
    response_path.write_text("", encoding="utf-8")
    run_id: str | None = None
    response = ""
    stderr = ""
    exit_code = 1
    interrupted = False

    try:
        begin_arguments = [
            "begin-run",
            args.session_id,
            "--agent",
            args.agent,
        ]
        if args.provider:
            begin_arguments.extend(["--provider", args.provider])
        if args.model:
            begin_arguments.extend(["--model", args.model])
        if args.lease_ttl_seconds:
            begin_arguments.extend(["--lease-ttl-seconds", str(args.lease_ttl_seconds)])
        begin = run_cli(args.home, *begin_arguments)
        run_id = begin["run_id"]
        run_cli(
            args.home,
            "append-user",
            run_id,
            "--session-id",
            args.session_id,
            "--input-file",
            str(prompt_path),
        )

        with response_path.open("wb") as response_handle, stderr_path.open("wb") as stderr_handle:
            process = subprocess.Popen(
                command,
                stdin=subprocess.PIPE,
                stdout=response_handle,
                stderr=stderr_handle,
            )
            try:
                process.communicate(input=prompt.encode("utf-8"))
            except KeyboardInterrupt:
                interrupted = True
                process.terminate()
                try:
                    process.wait(timeout=2)
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait()
            exit_code = process.returncode if process.returncode is not None else 1

        response = response_path.read_text(encoding="utf-8", errors="replace") if response_path.exists() else ""
        stderr = stderr_path.read_text(encoding="utf-8", errors="replace") if stderr_path.exists() else ""
        if not response and exit_code != 0:
            response = stderr or f"Agent command exited with status {exit_code}"
    except (OSError, RuntimeError, KeyError) as exc:
        exit_code = 1
        response = response or str(exc)
        stderr = stderr or str(exc)
    finally:
        if run_id:
            response_path.write_text(response, encoding="utf-8")
            assistant_status = "complete" if exit_code == 0 and not interrupted else "partial" if response else "failed"
            try:
                run_cli(
                    args.home,
                    "append-assistant",
                    run_id,
                    "--session-id",
                    args.session_id,
                    "--status",
                    assistant_status,
                    "--input-file",
                    str(response_path),
                )
            except (OSError, RuntimeError, KeyError):
                pass
            finish_status = "completed" if exit_code == 0 and not interrupted else "partial" if response else "failed"
            try:
                run_cli(
                    args.home,
                    "finish-run",
                    run_id,
                    "--session-id",
                    args.session_id,
                    "--status",
                    finish_status,
                )
            except (OSError, RuntimeError, KeyError):
                pass
        if stderr:
            sys.stderr.write(stderr)
        if response:
            sys.stdout.write(response)
        shutil.rmtree(temp_dir, ignore_errors=True)
    return 130 if interrupted else exit_code


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