#!/usr/bin/env python3
"""wsd — drive the desktop fleet over the homepage's realtime socket.

The primary channel. Unlike agent-desktop (needs a local cicy-code) and
hub-desktop (needs a hub login), this one has no precondition: every desktop
loads its homepage from the Worker and dials straight back, so it is reachable
whenever the app is running at all.

  wsd ls | nodes                          who is connected right now
  wsd exec <target> <shell command...>    run a shell command there
  wsd ipc  <target> <channel> [json-args] full IPC: ipcRenderer.invoke(channel,...)
  wsd team <target> <name>                declare which machine this is
  wsd rpc <target> <tool> [json-args]     call an Electron tool there
  wsd js  <target> <file|-|'code'>        run async JS in the page
  wsd main <target> <file|-|'code'>       run JS in the Electron MAIN process

<target> is a team (desktop-xs-1001, or just xs-1001), a hostname, a cid, or `all`.
Add --json anywhere for machine-readable output.
"""
import json, os, sys, urllib.request, urllib.error

CFG = os.path.expanduser("~/cicy-ai/db/desktop-ctrl.json")


def cfg():
    try:
        with open(CFG) as f:
            c = json.load(f)
        return c["base"].rstrip("/"), c["token"]
    except Exception as e:
        sys.exit(f"wsd: cannot read {CFG}: {e}")


def call(path, payload=None, timeout=130):
    base, tok = cfg()
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request(
        base + path, data=data,
        # Cloudflare's bot check answers urllib's default UA with 1010, so
        # present a normal one — this is our own origin, not a third party.
        headers={"x-cicy-ctrl": tok, "content-type": "application/json",
                 "user-agent": "Mozilla/5.0 (X11; Linux x86_64) wsd/1"},
        method="POST" if data else "GET")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return json.loads(r.read().decode())
    except urllib.error.HTTPError as e:
        body = e.read().decode()[:400]
        sys.exit(f"wsd: HTTP {e.code} {body}")
    except Exception as e:
        sys.exit(f"wsd: {e}")


def body(a):
    """A literal, a file path, or - for stdin."""
    if a == "-":
        return sys.stdin.read()
    return open(a).read() if os.path.exists(a) else a


JSON = False


def show(res):
    if JSON:
        print(json.dumps(res, ensure_ascii=False, indent=2)); return
    _show(res)


def _show(res):
    for r in res.get("results", []):
        out = r.get("out")
        if isinstance(out, str) and out and out[0] in "{[":
            try:
                out = json.dumps(json.loads(out), ensure_ascii=False)
            except Exception:
                pass
        print(f"{'ok ' if r.get('ok') else 'ERR'} {r.get('host','?'):<18} {out}")


def main():
    global JSON
    argv = [x for x in sys.argv[1:] if x != "--json"]
    JSON = "--json" in sys.argv
    if not argv or argv[0] in ("-h", "--help", "help"):
        sys.exit(__doc__)
    cmd, a = argv[0], argv[1:]

    if cmd in ("ls", "nodes"):
        d = call("/api/fleet")
        peers = d.get("peers", [])
        if JSON:
            print(json.dumps(peers, ensure_ascii=False, indent=2)); return
        print(f"build={d.get('build')}  peers={len(peers)}")
        for p in peers:
            who = p.get("team") or (p.get("host", "?") + " (no team)")
            print(f"  {who:<26} v={p.get('v') or '?':<9} {p.get('plat','?'):<6} "
                  f"auto={p.get('auto')} up={p.get('upSec')}s idle={p.get('idleSec')}s "
                  f"ip={p.get('ip')} cid={p.get('cid')}")
        return

    if cmd == "team":
        # Label a machine that predates the login gate. Identity is declared, so
        # this writes it where the page reads it and the socket reports it.
        if len(a) < 2:
            sys.exit("usage: wsd team <target> <name>")
        # Same namespacing the homepage applies: typed bare, stored prefixed.
        name = a[1] if a[1].startswith("desktop-") else "desktop-" + a[1]
        code = ('(()=>{const r=process.mainModule.require.bind(process.mainModule);'
                'const os=r("os"),fs=r("fs"),path=r("path");'
                'const p=path.join(os.homedir(),"cicy-ai","global.json");'
                'let c={};try{c=JSON.parse(fs.readFileSync(p,"utf8"))}catch(e){}'
                'c.desktopTeam=' + json.dumps(name) + ';'
                'fs.writeFileSync(p,JSON.stringify(c,null,2));'
                'return JSON.stringify({host:os.hostname(),team:c.desktopTeam})})()')
        wrapper = (
            "const g = await window.electronRPC('get_windows', {});"
            "const gt = ((g&&g.content)||[]).map(c=>c&&c.text).join('');"
            "let wid = 1; try { const ws = JSON.parse(gt); if (ws && ws.length) wid = ws[0].id; } catch (e) {}"
            "const r = await window.electronRPC('control_electron_BrowserWindow',"
            "{win_id: wid, code: " + json.dumps(code) + "});"
            "const out = ((r&&r.content)||[]).map(c=>c&&c.text).join('');"
            # Re-announce immediately; otherwise the label only lands on the next reconnect.
            "try { window.__cicyFleetRelabel && window.__cicyFleetRelabel(); } catch (e) {}"
            "return out;")
        return show(call("/api/rpc", {"target": a[0], "js": wrapper}))

    if cmd == "exec":
        if len(a) < 2:
            sys.exit("usage: wsd exec <target> <shell command...>")
        return show(call("/api/rpc", {"target": a[0], "tool": "exec_shell",
                                      "args": {"command": " ".join(a[1:])}}))

    if cmd == "ipc":
        # Full IPC over the id channel: wsd ipc <target> <channel> [json-args]
        if len(a) < 2:
            sys.exit("usage: wsd ipc <target> <channel> [json-args]")
        args = json.loads(a[2]) if len(a) > 2 else []
        return show(call("/api/rpc", {"target": a[0], "ipc": a[1], "args": args}))

    if cmd == "rpc":
        if len(a) < 2:
            sys.exit("usage: wsd rpc <target> <tool> [json-args]")
        args = json.loads(a[2]) if len(a) > 2 else {}
        return show(call("/api/rpc", {"target": a[0], "tool": a[1], "args": args}))

    if cmd == "js":
        if len(a) < 2:
            sys.exit("usage: wsd js <target> <file|-|code>")
        return show(call("/api/rpc", {"target": a[0], "js": body(a[1])}))

    if cmd == "main":
        if len(a) < 2:
            sys.exit("usage: wsd main <target> <file|-|code>")
        # Reaches the Electron MAIN process. The homepage bridge is the
        # unguarded one, so this is a plain call — no consent gate. Two things
        # bite here: there is no bare `require` in the eval sandbox (use
        # process.mainModule.require), and the window id is NOT 1 everywhere —
        # it is per-process, so resolve it instead of assuming.
        code = body(a[1])
        wrapper = (
            "const g = await window.electronRPC('get_windows', {});"
            "const gt = ((g&&g.content)||[]).map(c=>c&&c.text).join('');"
            "let wid = 1; try { const ws = JSON.parse(gt); if (ws && ws.length) wid = ws[0].id; } catch (e) {}"
            "const r = await window.electronRPC('control_electron_BrowserWindow',"
            "{win_id: wid, code: " + json.dumps(code) + "});"
            "return ((r&&r.content)||[]).map(c=>c&&c.text).join('');")
        return show(call("/api/rpc", {"target": a[0], "js": wrapper}))

    sys.exit(f"wsd: unknown command {cmd!r}")


if __name__ == "__main__":
    main()