#!/usr/bin/env python3
"""
helios-remote - CLI to control PCs connected to the Helios Remote Relay Server.
Config is loaded from config.env in the same directory as this script.
"""

import argparse
import base64
import os
import sys
from pathlib import Path

# ── Load config.env ──────────────────────────────────────────────────────────
_config_path = Path(__file__).parent / "config.env"
_config: dict = {}
if _config_path.exists():
    for _line in _config_path.read_text().splitlines():
        _line = _line.strip()
        if _line and not _line.startswith("#") and "=" in _line:
            k, v = _line.split("=", 1)
            _config[k.strip()] = v.strip()

BASE_URL = _config.get("HELIOS_REMOTE_URL", "").rstrip("/")
API_KEY  = _config.get("HELIOS_REMOTE_API_KEY", "")

if not BASE_URL or not API_KEY:
    sys.exit(f"[helios-remote] ERROR: config.env missing or incomplete at {_config_path}")

try:
    import requests
except ImportError:
    sys.exit("[helios-remote] ERROR: 'requests' not installed. Run: pip install requests")


# ── HTTP helpers ─────────────────────────────────────────────────────────────

def _headers() -> dict:
    return {"X-Api-Key": API_KEY, "Content-Type": "application/json"}


def _req(method: str, path: str, timeout: int = 30, **kwargs):
    url = f"{BASE_URL}{path}"
    try:
        resp = requests.request(method, url, headers=_headers(), timeout=timeout, **kwargs)
    except requests.exceptions.ConnectionError as exc:
        sys.exit(f"[helios-remote] CONNECTION ERROR: Cannot reach {url}\n  → {exc}")
    except requests.exceptions.Timeout:
        sys.exit(f"[helios-remote] TIMEOUT: {url} did not respond within {timeout} s")

    if not resp.ok:
        body = resp.text[:1000] if resp.text else "(empty body)"
        sys.exit(
            f"[helios-remote] HTTP {resp.status_code} {resp.reason}\n"
            f"  URL    : {url}\n"
            f"  Method : {method.upper()}\n"
            f"  Body   : {body}"
        )
    return resp


# ── Subcommands ──────────────────────────────────────────────────────────────

HEADERS = _headers  # alias for lazy evaluation


def resolve_session(session_id: str) -> str:
    """If session_id looks like a label (not a UUID), resolve it to an actual UUID."""
    import re
    uuid_pattern = re.compile(
        r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$', re.I
    )
    if uuid_pattern.match(session_id):
        return session_id
    # Look up by label
    resp = _req("GET", "/sessions")
    sessions = resp.json().get("sessions", [])
    for s in sessions:
        if s.get("label", "").lower() == session_id.lower():
            return s["id"]
    raise SystemExit(f"[helios-remote] No session found with label '{session_id}'")


def resolve_window(sid: str, window_id_or_name: str) -> int:
    """If window_id_or_name is a number, return it. Otherwise search by title substring."""
    if window_id_or_name.lstrip('-').isdigit():
        return int(window_id_or_name)
    # Search by title
    resp = _req("GET", f"/sessions/{sid}/windows")
    windows = resp.json().get("windows", [])
    query = window_id_or_name.lower()
    matches = [w for w in windows if w.get("visible") and query in w.get("title", "").lower()]
    if not matches:
        raise SystemExit(f"[helios-remote] No visible window matching '{window_id_or_name}'")
    if len(matches) > 1:
        print(f"[helios-remote] Multiple matches for '{window_id_or_name}', using first:")
        for w in matches:
            print(f"  {w['id']}  {w['title']}")
    return int(matches[0]["id"])


def cmd_sessions(_args):
    """List all connected sessions."""
    resp = _req("GET", "/sessions")
    data = resp.json()
    sessions = data if isinstance(data, list) else data.get("sessions", [])

    if not sessions:
        print("No sessions connected.")
        return

    print(f"{'ID':<36}  {'Label':<20}  Info")
    print("-" * 70)
    for s in sessions:
        sid   = str(s.get("id") or s.get("session_id") or "?")
        label = str(s.get("label") or s.get("name") or "")
        info  = str(s.get("os") or s.get("platform") or s.get("info") or "")
        print(f"{sid:<36}  {label:<20}  {info}")


def cmd_label(args):
    """Assign a human-readable name to a session."""
    sid = resolve_session(args.session_id)
    _req("POST", f"/sessions/{sid}/label", json={"label": args.name})
    print(f"Session {sid!r} labeled as {args.name!r}.")


def cmd_screenshot(args):
    """Capture a screenshot → /tmp/helios-remote-screenshot.png"""
    sid = resolve_session(args.session_id)
    resp = _req("POST", f"/sessions/{sid}/screenshot")
    out_path = Path("/tmp/helios-remote-screenshot.png")

    content_type = resp.headers.get("Content-Type", "")
    if "image" in content_type:
        out_path.write_bytes(resp.content)
    else:
        data = resp.json()
        b64 = (
            data.get("screenshot")
            or data.get("image")
            or data.get("image_base64")
            or data.get("data")
            or data.get("png")
        )
        if not b64:
            sys.exit(
                f"[helios-remote] ERROR: Screenshot response has no image field.\n"
                f"  Content-Type : {content_type}\n"
                f"  Keys present : {list(data.keys()) if isinstance(data, dict) else type(data).__name__}"
            )
        if "," in b64:          # strip data-URI prefix (data:image/png;base64,...)
            b64 = b64.split(",", 1)[1]
        out_path.write_bytes(base64.b64decode(b64))

    print(str(out_path))


def cmd_exec(args):
    """Run a shell command on the remote session."""
    sid = resolve_session(args.session_id)
    command = " ".join(args.parts) if isinstance(args.parts, list) else args.parts
    body = {"command": command}
    if args.timeout:
        body["timeout_ms"] = args.timeout * 1000  # seconds → ms
    resp = _req("POST", f"/sessions/{sid}/exec", json=body, timeout=max(35, (args.timeout or 30) + 5))
    data = resp.json()

    stdout    = data.get("stdout") or data.get("output") or ""
    stderr    = data.get("stderr") or ""
    exit_code = data.get("exit_code")

    if stdout:
        print(stdout, end="" if stdout.endswith("\n") else "\n")
    if stderr:
        print("[stderr]", stderr, file=sys.stderr)
    if exit_code is not None and int(exit_code) != 0:
        print(f"[helios-remote] Command exited with code {exit_code}", file=sys.stderr)
        sys.exit(int(exit_code))


def cmd_click(args):
    """Send a mouse click to the remote session."""
    sid = resolve_session(args.session_id)
    _req("POST", f"/sessions/{sid}/click", json={"x": args.x, "y": args.y})
    print(f"Clicked ({args.x}, {args.y}) on session {sid!r}.")


def cmd_type(args):
    """Send keyboard input to the remote session."""
    sid = resolve_session(args.session_id)
    _req("POST", f"/sessions/{sid}/type", json={"text": args.text})
    print(f"Typed {len(args.text)} chars on session {sid!r}.")


def cmd_windows(args):
    """List windows on the remote session (visible only by default)."""
    sid = resolve_session(args.session_id)
    resp = _req("GET", f"/sessions/{sid}/windows")
    data = resp.json()
    windows = data.get("windows", [])
    # Filter to visible only unless --all is passed
    all_windows = getattr(args, "all", False)
    if not all_windows:
        windows = [w for w in windows if w.get("visible")]
    if not windows:
        print("No windows returned.")
        return
    print(f"{'ID':<20}  Title")
    print("-" * 70)
    for w in windows:
        wid   = str(w.get("id", "?"))
        title = w.get("title", "")
        print(f"{wid:<20}  {title}")


def cmd_minimize_all(args):
    """Minimize all windows on the remote session."""
    sid = resolve_session(args.session_id)
    _req("POST", f"/sessions/{sid}/windows/minimize-all")
    print(f"All windows minimized on session {sid!r}.")


def cmd_focus(args):
    """Bring a window to the foreground (by ID or title substring)."""
    sid = resolve_session(args.session_id)
    wid = resolve_window(sid, args.window_id)
    _req("POST", f"/sessions/{sid}/windows/{wid}/focus")
    print(f"Window {wid} focused on session {sid!r}.")


def cmd_maximize(args):
    """Maximize and focus a window (by ID or title substring)."""
    sid = resolve_session(args.session_id)
    wid = resolve_window(sid, args.window_id)
    _req("POST", f"/sessions/{sid}/windows/{wid}/maximize")
    print(f"Window {wid} maximized on session {sid!r}.")


def cmd_logs(args):
    """Fetch the last N lines of the client log file."""
    sid = resolve_session(args.session_id)
    resp = _req("GET", f"/sessions/{sid}/logs", params={"lines": args.lines})
    data = resp.json()
    if "error" in data:
        sys.exit(f"[helios-remote] {data['error']}")
    print(f"# {data.get('log_path', '?')}  (last {args.lines} lines)")
    print(data.get("content", ""))


def cmd_screenshot_window(args):
    """Capture a specific window by ID or title substring → /tmp/helios-remote-screenshot.png"""
    sid = resolve_session(args.session_id)
    wid = resolve_window(sid, args.window_id)
    resp = _req("POST", f"/sessions/{sid}/windows/{wid}/screenshot")
    data = resp.json()
    if "error" in data:
        sys.exit(f"[helios-remote] {data['error']}")
    import base64, os
    out_path = args.output or "/tmp/helios-remote-screenshot.png"
    img_bytes = base64.b64decode(data["image_base64"])
    with open(out_path, "wb") as f:
        f.write(img_bytes)
    print(out_path)
    return out_path


def _script_commit() -> str:
    """Return the git commit hash of remote.py itself."""
    import subprocess, os
    try:
        r = subprocess.run(
            ["git", "log", "-1", "--format=%h", "--", __file__],
            capture_output=True, text=True,
            cwd=os.path.dirname(os.path.abspath(__file__))
        )
        return r.stdout.strip() or "unknown"
    except Exception:
        return "unknown"


def cmd_status(args):
    """Show relay, remote.py, and client commit in one call."""
    import requests as _requests

    # 1. Relay (public endpoint, no auth)
    try:
        r = _requests.get(f"{BASE_URL}/version", timeout=10)
        relay = r.json() if r.ok else {"commit": f"HTTP {r.status_code}"}
    except Exception as e:
        relay = {"commit": str(e)}

    # 2. Client (via session)
    sid = resolve_session(args.session_id)
    try:
        r2 = _req("GET", f"/sessions/{sid}/version")
        client = r2.json()
    except Exception as e:
        client = {"commit": str(e)}

    relay_commit  = relay.get("commit", "?")
    script_commit = _script_commit()
    client_commit = client.get("commit", "?")
    all_same = relay_commit == script_commit == client_commit

    print(f"  relay      {relay_commit}")
    print(f"  remote.py  {script_commit}")
    print(f"  client     {client_commit}")
    print(f"  {'✅ all in sync' if all_same else '⚠️  OUT OF SYNC'}")


def cmd_server_version(_args):
    """Get server version (no auth required)."""
    import requests as _requests
    try:
        resp = _requests.get(f"{BASE_URL}/version", timeout=10)
    except Exception as exc:
        sys.exit(f"[helios-remote] CONNECTION ERROR: {exc}")
    if not resp.ok:
        sys.exit(f"[helios-remote] HTTP {resp.status_code}: {resp.text[:500]}")
    data = resp.json()
    print(f"Server version : {data.get('version', '?')}")
    print(f"Commit         : {data.get('commit', '?')}")


def cmd_version(args):
    """Get the client version for a session."""
    sid = resolve_session(args.session_id)
    resp = _req("GET", f"/sessions/{sid}/version")
    data = resp.json()
    print(f"Client version : {data.get('version', '?')}")
    print(f"Commit         : {data.get('commit', '?')}")


def cmd_upload(args):
    """Upload a local file to the remote session."""
    sid = resolve_session(args.session_id)
    local_path = Path(args.local_path)
    if not local_path.exists():
        sys.exit(f"[helios-remote] ERROR: Local file not found: {local_path}")
    content_base64 = base64.b64encode(local_path.read_bytes()).decode()
    _req("POST", f"/sessions/{sid}/upload", json={
        "path": args.remote_path,
        "content_base64": content_base64,
    })
    print(f"Uploaded {local_path} → {args.remote_path} on session {sid!r}.")


def cmd_download(args):
    """Download a file from the remote session to a local path."""
    sid = resolve_session(args.session_id)
    from urllib.parse import quote
    encoded_path = quote(args.remote_path, safe="")
    resp = _req("GET", f"/sessions/{sid}/download?path={encoded_path}")
    data = resp.json()
    b64 = data.get("content_base64", "")
    if not b64:
        sys.exit(f"[helios-remote] ERROR: No content in download response. Keys: {list(data.keys())}")
    local_path = Path(args.local_path)
    local_path.parent.mkdir(parents=True, exist_ok=True)
    local_path.write_bytes(base64.b64decode(b64))
    size = data.get("size", len(base64.b64decode(b64)))
    print(f"Downloaded {args.remote_path} → {local_path} ({size} bytes).")


def cmd_find_window(args):
    """Find visible windows by title substring on the remote session."""
    sid = resolve_session(args.session_id)
    resp = _req("GET", f"/sessions/{sid}/windows")
    data = resp.json()
    windows = data.get("windows", [])
    query = args.title.lower()
    # Only visible windows
    matches = [w for w in windows if w.get("visible") and query in w.get("title", "").lower()]
    if not matches:
        print(f"No visible windows matching '{args.title}'.")
        return
    print(f"{'ID':<20}  Title")
    print("-" * 70)
    for w in matches:
        wid   = str(w.get("id", "?"))
        title = w.get("title", "")
        print(f"{wid:<20}  {title}")


def cmd_prompt(args):
    """Show a MessageBox on the remote PC asking the user to do something.
    Blocks until the user clicks OK — use this when the AI needs the user
    to perform a manual action (e.g. click a button, confirm a dialog)."""
    sid = resolve_session(args.session_id)
    body = {"message": args.message}
    if args.title:
        body["title"] = args.title
    _req("POST", f"/sessions/{sid}/prompt", json=body)
    print(f"User confirmed prompt on session {sid!r}.")


def cmd_wait_for_window(args):
    """Poll until a visible window with the given title appears (or timeout)."""
    import time
    sid = resolve_session(args.session_id)
    query = args.title.lower()
    deadline = time.time() + args.timeout
    interval = 1.0
    print(f"Waiting for window matching '{args.title}' (timeout: {args.timeout}s)...")
    while time.time() < deadline:
        resp = _req("GET", f"/sessions/{sid}/windows")
        windows = resp.json().get("windows", [])
        matches = [w for w in windows if w.get("visible") and query in w.get("title", "").lower()]
        if matches:
            print(f"{'ID':<20}  Title")
            print("-" * 70)
            for w in matches:
                print(f"{str(w.get('id','?')):<20}  {w.get('title','')}")
            return
        time.sleep(interval)
    sys.exit(f"[helios-remote] Timeout: no window matching '{args.title}' appeared within {args.timeout}s")


def cmd_run(args):
    """Launch a program on the remote session (fire-and-forget)."""
    sid = resolve_session(args.session_id)
    _req("POST", f"/sessions/{sid}/run", json={"program": args.program, "args": args.args})
    print(f"Started {args.program!r} on session {sid!r}.")


def cmd_clipboard_get(args):
    """Get the clipboard contents from the remote session."""
    sid = resolve_session(args.session_id)
    resp = _req("GET", f"/sessions/{sid}/clipboard")
    data = resp.json()
    print(data.get("text", ""))


def cmd_clipboard_set(args):
    """Set the clipboard contents on the remote session."""
    sid = resolve_session(args.session_id)
    _req("POST", f"/sessions/{sid}/clipboard", json={"text": args.text})
    print(f"Clipboard set ({len(args.text)} chars) on session {sid!r}.")


# ── CLI wiring ────────────────────────────────────────────────────────────────

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="remote.py",
        description="Control PCs connected to the Helios Remote Relay Server.",
    )
    sub = p.add_subparsers(dest="subcmd", required=True)

    sub.add_parser("sessions", help="List all connected sessions")

    lp = sub.add_parser("label", help="Assign a name to a session")
    lp.add_argument("session_id")
    lp.add_argument("name")

    sp = sub.add_parser("screenshot", help="Capture screenshot → /tmp/helios-remote-screenshot.png")
    sp.add_argument("session_id")

    swp = sub.add_parser("screenshot-window", help="Capture a specific window (by ID or title)")
    swp.add_argument("session_id")
    swp.add_argument("window_id", help="Window ID (number) or title substring")
    swp.add_argument("--output", default=None, help="Output path (default: /tmp/helios-remote-screenshot.png)")
    swp.set_defaults(func=cmd_screenshot_window)

    ep = sub.add_parser("exec", help="Run a shell command on the remote session")
    ep.add_argument("session_id")
    ep.add_argument("parts", nargs=argparse.REMAINDER, metavar="command",
                    help="Command (and arguments) to execute")
    ep.add_argument("--timeout", type=int, default=None, metavar="SECONDS",
                    help="Timeout in seconds (default: 30). Use higher for long downloads etc.")

    cp = sub.add_parser("click", help="Send a mouse click")
    cp.add_argument("session_id")
    cp.add_argument("x", type=int)
    cp.add_argument("y", type=int)

    tp = sub.add_parser("type", help="Send keyboard input")
    tp.add_argument("session_id")
    tp.add_argument("text")

    wp = sub.add_parser("windows", help="List all windows on the remote session")
    wp.add_argument("session_id")

    mp = sub.add_parser("minimize-all", help="Minimize all windows")
    mp.add_argument("session_id")

    fp = sub.add_parser("focus", help="Bring a window to the foreground")
    fp.add_argument("session_id")
    fp.add_argument("window_id")

    xp = sub.add_parser("maximize", help="Maximize and focus a window")
    xp.add_argument("session_id")
    xp.add_argument("window_id")

    stp = sub.add_parser("status", help="Show relay + client commit and sync status")
    stp.add_argument("session_id")

    lgp = sub.add_parser("logs", help="Fetch last N lines of client log file")
    lgp.add_argument("session_id")
    lgp.add_argument("--lines", type=int, default=100, metavar="N", help="Number of lines (default: 100)")

    sub.add_parser("server-version", help="Get server version (no auth required)")

    vp = sub.add_parser("version", help="Get client version for a session")
    vp.add_argument("session_id")

    up = sub.add_parser("upload", help="Upload a local file to the remote session")
    up.add_argument("session_id")
    up.add_argument("local_path")
    up.add_argument("remote_path")

    dp = sub.add_parser("download", help="Download a file from the remote session")
    dp.add_argument("session_id")
    dp.add_argument("remote_path")
    dp.add_argument("local_path")

    fwp = sub.add_parser("find-window", help="Find windows by title substring")
    fwp.add_argument("session_id")
    fwp.add_argument("title", help="Substring to search for (case-insensitive)")

    wfwp = sub.add_parser("wait-for-window", help="Poll until a window with given title appears")
    wfwp.add_argument("session_id")
    wfwp.add_argument("title", help="Substring to wait for (case-insensitive)")
    wfwp.add_argument("--timeout", type=int, default=30, metavar="SECONDS", help="Max wait time (default: 30s)")
    wfwp.set_defaults(func=cmd_wait_for_window)

    pp = sub.add_parser("prompt", help="Show a MessageBox asking the user to do something manually")
    pp.add_argument("session_id")
    pp.add_argument("message", help="What to ask the user (e.g. 'Please click Save, then OK')")
    pp.add_argument("--title", default=None, help="Dialog title (default: Helios Remote)")
    pp.set_defaults(func=cmd_prompt)

    rp = sub.add_parser("run", help="Launch a program on the remote session (fire-and-forget)")
    rp.add_argument("session_id")
    rp.add_argument("program", help="Program to launch (e.g. notepad.exe)")
    rp.add_argument("args", nargs=argparse.REMAINDER, help="Optional arguments")

    cgp = sub.add_parser("clipboard-get", help="Get clipboard contents from the remote session")
    cgp.add_argument("session_id")

    csp = sub.add_parser("clipboard-set", help="Set clipboard contents on the remote session")
    csp.add_argument("session_id")
    csp.add_argument("text")

    return p


def main():
    parser = build_parser()
    args = parser.parse_args()

    {
        "sessions":     cmd_sessions,
        "label":        cmd_label,
        "screenshot":   cmd_screenshot,
        "exec":         cmd_exec,
        "click":        cmd_click,
        "type":         cmd_type,
        "windows":      cmd_windows,
        "minimize-all": cmd_minimize_all,
        "focus":        cmd_focus,
        "maximize":       cmd_maximize,
        "server-version": cmd_server_version,
        "version":        cmd_version,
        "upload":         cmd_upload,
        "download":       cmd_download,
        "status":            cmd_status,
        "logs":              cmd_logs,
        "screenshot-window": cmd_screenshot_window,
        "find-window":      cmd_find_window,
        "wait-for-window":  cmd_wait_for_window,
        "run":              cmd_run,
        "prompt":           cmd_prompt,
        "clipboard-get":    cmd_clipboard_get,
        "clipboard-set":    cmd_clipboard_set,
    }[args.subcmd](args)


if __name__ == "__main__":
    main()