#!/usr/bin/env python3
# Fetch a file from the guest via the QEMU guest agent protocol
# (/tmp/qemu-guest-agent.sock, the virtio-serial QGA chardev - the same
# socket vgpu_vm.py exec uses). Usage: qga_fetch.py <guest_path> <host_path> [chunk]
#
# Robustness (r1490): the r1489 evidence capture fetched nothing because the
# guest-side frame was still being written / the agent answered late. Retry
# the open a few times, tolerate error replies with a clear diagnostic, use a
# long socket timeout, and write to a temporary file so a partial transfer is
# never mistaken for a complete capture.
import base64
import json
import os
import socket
import sys
import time

if len(sys.argv) < 3:
    print("usage: qga_fetch.py <guest_path> <host_path> [chunk]", file=sys.stderr)
    sys.exit(2)

GUEST = sys.argv[1]
HOST = sys.argv[2]
CHUNK = int(sys.argv[3]) if len(sys.argv) > 3 else 1 << 20
SOCK = "/tmp/qemu-guest-agent.sock"
TRIES = 6
TMP = HOST + ".part"


def recv_line(s):
    buf = b""
    while b"\n" not in buf:
        chunk = s.recv(65536)
        if not chunk:
            break
        buf += chunk
    return buf.split(b"\n", 1)[0]


def cmd(s, obj):
    s.sendall((json.dumps(obj) + "\r\n").encode())
    line = recv_line(s)
    return json.loads(line)


def open_guest():
    last = None
    for i in range(TRIES):
        try:
            s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            s.settimeout(120.0)
            s.connect(SOCK)
            r = cmd(s, {"execute": "guest-file-open",
                        "arguments": {"path": GUEST, "mode": "r"}})
            if "return" in r:
                return s, r["return"], None
            last = f"guest-file-open error: {r.get('error')}"
        except Exception as e:  # noqa: BLE001 - diagnostic surface
            last = f"guest-file-open exception: {e}"
        try:
            s.close()
        except Exception:
            pass
        time.sleep(1.5)
        if i + 1 < TRIES:
            print(f"retry {i + 1}/{TRIES}: {last}", file=sys.stderr)
    return None, None, last


s, h, err = open_guest()
if s is None:
    print(f"FETCH-FAIL {GUEST} -> {HOST}: {err}", file=sys.stderr)
    sys.exit(1)

print(f"handle={h}")
total = 0
try:
    with open(TMP, "wb") as out:
        while True:
            r = cmd(s, {"execute": "guest-file-read",
                        "arguments": {"handle": h, "count": CHUNK}})
            if "return" not in r:
                raise RuntimeError(f"guest-file-read error: {r.get('error')}")
            d = r["return"]
            buf = base64.b64decode(d.get("buf-b64", ""))
            if buf:
                out.write(buf)
                total += len(buf)
                print(f"\r{total}", end="", flush=True)
            if d.get("eof"):
                break
    cmd(s, {"execute": "guest-file-close", "arguments": {"handle": h}})
except Exception as e:  # noqa: BLE001
    print(f"\nFETCH-FAIL {GUEST} -> {HOST}: {e}", file=sys.stderr)
    try:
        if os.path.exists(TMP):
            os.unlink(TMP)
    except OSError:
        pass
    sys.exit(1)
finally:
    try:
        s.close()
    except Exception:
        pass

os.replace(TMP, HOST)
print(f"\nFETCHED {total} bytes -> {HOST}")
sys.exit(0)