#!/usr/bin/env python3
"""Drive a VM session by hand, with no model involved.

    sudo -u owuivmd vmctl open --os ubuntu --project demo
    sudo -u owuivmd vmctl plan "write a script" "run it"
    sudo -u owuivmd vmctl run 'uname -a'
    sudo -u owuivmd vmctl shot /tmp/screen.png
    sudo -u owuivmd vmctl close

This exists because the agent loop should be the LAST thing added, not
the first. Every verb the sub-agent will eventually call is reachable
here in the same order, so a session that misbehaves can be reproduced
without an LLM in the middle deciding what to try next - which is the
difference between debugging a VM and debugging a conversation.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys

# Installed, vmctl lives in /usr/local/bin while the modules live in
# /usr/local/lib/owuivm; in the repo they sit side by side. Both are
# tried so the same file works from either place.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, "/usr/local/lib/owuivm")
from client import Vmd, VmdError, VmdUnavailable  # noqa: E402


def main() -> int:
    parser = argparse.ArgumentParser(prog="vmctl", description=__doc__)
    sub = parser.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("open"); p.add_argument("--os", default="ubuntu"); p.add_argument("--project")
    sub.add_parser("close"); sub.add_parser("status"); sub.add_parser("cancel")
    sub.add_parser("reap"); sub.add_parser("report"); sub.add_parser("internet")
    # WITHOUT THIS, DRIVING A SESSION BY HAND DOES NOT WORK.
    #
    # `ready` is set by the session_ready verb and by nothing else, and
    # every other verb that touches the machine is refused until it is.
    # The agent loop calls it as a matter of course; an operator at a
    # terminal had no way to, so `open` then `run` came back "the machine
    # is still booting" for as long as you cared to wait - with the guest
    # up and answering the whole time.
    p = sub.add_parser("ready", help="wait for the machine to finish booting")
    p.add_argument("--timeout", type=float, default=300)
    p.add_argument("--desktop", action="store_true",
                   help="also wait for the desktop session, for GUI work")
    p = sub.add_parser("plan"); p.add_argument("tasks", nargs="+")
    p = sub.add_parser("task"); p.add_argument("id", type=int); p.add_argument("status"); p.add_argument("note", nargs="?", default="")
    p = sub.add_parser("run"); p.add_argument("command"); p.add_argument("--timeout", type=float, default=120)
    p = sub.add_parser("write"); p.add_argument("path"); p.add_argument("content")
    p = sub.add_parser("read"); p.add_argument("path")
    p = sub.add_parser("ls"); p.add_argument("path", nargs="?", default=".")
    p = sub.add_parser("shot"); p.add_argument("out")
    p = sub.add_parser("click"); p.add_argument("x", type=int); p.add_argument("y", type=int)
    p = sub.add_parser("key"); p.add_argument("keys")
    p = sub.add_parser("type"); p.add_argument("text")
    p = sub.add_parser("commit"); p.add_argument("message")
    p = sub.add_parser("finish"); p.add_argument("summary"); p.add_argument("--check", required=True); p.add_argument("--expect", default="")
    p = sub.add_parser("snapshot"); p.add_argument("--os", required=True)
    p.add_argument("--settle", type=float, default=45.0)
    sub.add_parser("snapshots")
    p = sub.add_parser("raw"); p.add_argument("verb"); p.add_argument("json", nargs="?", default="{}")
    # Reading the record needs no daemon and no machine: these are files.
    sub.add_parser("logs", help="list the recorded sessions, newest last")
    p = sub.add_parser("log", help="print one recorded session in full")
    p.add_argument("which", nargs="?", default="",
                   help="a file name, a session id, a chat id, or nothing for the newest")
    p.add_argument("--json", action="store_true", help="the raw records instead")

    args = parser.parse_args()

    # BEFORE Vmd(): the whole point of these is to work when the session
    # is over, the machine is gone and the daemon has nothing to say.
    if args.cmd in ("logs", "log"):
        return _logs(args)

    try:
        with Vmd() as vmd:
            result = dispatch(vmd, args)
    except VmdUnavailable as exc:
        print(f"vmctl: {exc}", file=sys.stderr)
        return 69
    except VmdError as exc:
        # A refusal is an answer, not a crash: these are the same strings
        # the sub-agent will read, so they are printed plainly to be judged.
        print(f"refused: {exc}", file=sys.stderr)
        return 1

    if isinstance(result, str):
        print(result)
    else:
        print(json.dumps(result, indent=2)[:20000])
    return 0


def _logs(a: argparse.Namespace) -> int:
    """Read the session logs back. See sessionlog.py for why they exist."""
    import sessionlog

    found = sessionlog.sessions()
    if not found:
        print(f"no sessions recorded yet in {sessionlog.LOG_DIR}", file=sys.stderr)
        return 1

    if a.cmd == "logs":
        print(f"{'when':<20} {'os':<8} {'project':<26} {'took':>7}  outcome")
        for entry in found:
            took = entry.get("seconds")
            took = f"{int(took // 60)}m{int(took % 60):02d}s" if took else "-"
            print(f"{entry['when']:<20} {entry['os']:<8} {entry['project'][:26]:<26} "
                  f"{took:>7}  {entry['outcome']}")
            print(f"    {entry['file']}   chat {entry['chat'] or '-'}")
        return 0

    want = a.which.strip()
    if not want:
        chosen = found[-1]
    else:
        matches = [e for e in found
                   if want in e["file"] or want in e["session"]
                   or (e["chat"] and want in e["chat"]) or want in e["project"]]
        if not matches:
            print(f"no recorded session matches {want!r}", file=sys.stderr)
            return 1
        chosen = matches[-1]

    if a.json:
        for record in sessionlog.load(chosen["path"]):
            print(json.dumps(record))
        return 0
    print(sessionlog.render(chosen["path"]))
    return 0


def dispatch(vmd: Vmd, a: argparse.Namespace):
    match a.cmd:
        case "open":     return vmd.call("session_open", os=a.os, project=a.project, owner="vmctl")
        case "close":    return vmd.call("session_close")
        case "status":   return vmd.call("session_status")
        case "cancel":   return vmd.call("session_cancel")
        case "reap":     return vmd.call("reap")
        case "report":   return vmd.call("report")["text"]
        case "internet": return vmd.call("check_internet")
        case "ready":    return vmd.call("session_ready", timeout=a.timeout,
                                         desktop=a.desktop)
        case "plan":     return vmd.call("plan_tasks", tasks=a.tasks)["tasks"]
        case "task":     return vmd.call("update_task", id=a.id, status=a.status, note=a.note)["tasks"]
        case "run":      return vmd.call("run", command=a.command, timeout_s=a.timeout)
        case "write":    return vmd.call("write_file", path=a.path, content=a.content)
        case "read":     return vmd.call("read_file", path=a.path)
        case "ls":       return vmd.call("list_dir", path=a.path)["listing"]
        case "click":    return vmd.call("click", x=a.x, y=a.y)
        case "key":      return vmd.call("key", keys=a.keys)
        case "type":     return vmd.call("type", text=a.text)
        case "commit":   return vmd.call("git_commit", message=a.message)
        case "finish":
            return vmd.call("finish", summary=a.summary,
                            verification={"command": a.check, "expected": a.expect})
        case "shot":
            shot = vmd.call("screenshot")
            with open(a.out, "wb") as handle:
                handle.write(base64.b64decode(shot["png_b64"]))
            return {"written": a.out, "width": shot["width"], "height": shot["height"]}
        case "snapshot":
            print(json.dumps(vmd.call("snapshot", os=a.os, settle=a.settle), indent=2))
        case "snapshots":
            print(json.dumps(vmd.call("snapshot_status"), indent=2))
        case "raw":      return vmd.call(a.verb, **json.loads(a.json))
    raise SystemExit(2)


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