#!/usr/bin/env python3
"""AirXDB remote GUI device helper over SSH."""
from __future__ import annotations
import argparse, base64, json, os, shlex, shutil, subprocess
from datetime import datetime, timezone
from pathlib import Path

def stamp(): return datetime.now().strftime('%Y%m%d-%H%M%S')
def iso(): return datetime.now(timezone.utc).isoformat()
def q(s): return shlex.quote(str(s))

def read_env(path):
    vals = {}
    if path.exists():
        for raw in path.read_text(encoding='utf-8-sig').splitlines():
            line = raw.strip()
            if line and not line.startswith('#') and '=' in line:
                k, v = line.split('=', 1); vals[k.strip()] = v.strip().strip('"').strip("'")
    for k in ['AIRXDB_REMOTE_SSH_TARGET','AIRXDB_REMOTE_SSH_PORT','AIRXDB_REMOTE_SSH_OPTIONS','AIRXDB_REMOTE_WORKDIR','AIRXDB_REMOTE_SCREENSHOT_TOOL','AIRXDB_REMOTE_DISPLAY']:
        if os.environ.get(k): vals[k] = os.environ[k]
    return vals

def write_env(path, updates):
    path.parent.mkdir(parents=True, exist_ok=True)
    old = path.read_text(encoding='utf-8') if path.exists() else ''
    keys = set(updates); lines = []
    for raw in old.splitlines():
        key = raw.split('=',1)[0].strip() if '=' in raw and not raw.strip().startswith('#') else None
        if key not in keys: lines.append(raw)
    if lines and lines[-1].strip(): lines.append('')
    for k, v in updates.items():
        if v: lines.append(f'{k}={v}')
    path.write_text('\n'.join(lines).rstrip()+'\n', encoding='utf-8', newline='\n')

def ensure_files(project):
    d = project/'AirPlan'/'state'/'airxdb'; d.mkdir(parents=True, exist_ok=True)
    ex = d/'remote-device.env.example'
    if not ex.exists():
        ex.write_text('# AirXDB remote GUI device configuration\nAIRXDB_REMOTE_SSH_TARGET=user@host\nAIRXDB_REMOTE_SSH_PORT=22\nAIRXDB_REMOTE_SSH_OPTIONS=\nAIRXDB_REMOTE_WORKDIR=\nAIRXDB_REMOTE_SCREENSHOT_TOOL=auto\nAIRXDB_REMOTE_DISPLAY=\n', encoding='utf-8', newline='\n')
    gi = d/'.gitignore'; old = gi.read_text(encoding='utf-8') if gi.exists() else ''
    for item in ['remote-device.env','midscene.local.env']:
        if item not in old.splitlines(): old = (old.rstrip()+f'\n{item}\n').lstrip()
    gi.write_text(old, encoding='utf-8', newline='\n')

def ssh_args(env, cmd):
    target = env.get('AIRXDB_REMOTE_SSH_TARGET','')
    if not target: raise SystemExit('AIRXDB_REMOTE_SSH_TARGET is required')
    args = [shutil.which('ssh') or 'ssh']
    port = env.get('AIRXDB_REMOTE_SSH_PORT','22')
    if port: args += ['-p', port]
    opts = env.get('AIRXDB_REMOTE_SSH_OPTIONS','').strip()
    if opts: args += shlex.split(opts, posix=True)
    return args + [target, cmd]

def run(env, cmd, timeout=60):
    return subprocess.run(ssh_args(env, cmd), capture_output=True, text=True, encoding='utf-8', errors='replace', timeout=timeout)

def remote_home(env):
    r = run(env, 'printf %s "$HOME"', 15)
    return r.stdout.strip() if r.returncode == 0 else ''

def workdir(env):
    if env.get('AIRXDB_REMOTE_WORKDIR'): return env['AIRXDB_REMOTE_WORKDIR']
    home = remote_home(env); return (home.rstrip('/')+'/.airxdb') if home else '.airxdb'

def detect_tool(env):
    tool = env.get('AIRXDB_REMOTE_SCREENSHOT_TOOL','auto').strip()
    if tool and tool != 'auto':
        r = run(env, f'command -v {q(tool)}', 15); return tool if r.returncode == 0 else ''
    script = 'for t in gnome-screenshot spectacle scrot grim import screencapture; do command -v "$t" >/dev/null 2>&1 && { printf %s "$t"; exit 0; }; done'
    r = run(env, script, 15)
    return r.stdout.strip().splitlines()[-1] if r.returncode == 0 and r.stdout.strip() else ''

def install_tool(env):
    script = """set -e
if command -v scrot >/dev/null 2>&1; then exit 0; fi
if [ "$(id -u 2>/dev/null || echo 1)" = "0" ]; then SUDO=""; else SUDO="sudo -n"; fi
if command -v apt-get >/dev/null 2>&1; then $SUDO apt-get update && $SUDO apt-get install -y scrot;
elif command -v dnf >/dev/null 2>&1; then $SUDO dnf install -y scrot;
elif command -v yum >/dev/null 2>&1; then $SUDO yum install -y scrot;
elif command -v apk >/dev/null 2>&1; then $SUDO apk add scrot;
elif command -v pacman >/dev/null 2>&1; then $SUDO pacman -Sy --noconfirm scrot;
else exit 42; fi"""
    return run(env, script, 300)

def setup(project, env, auto=True):
    if not env.get('AIRXDB_REMOTE_SSH_TARGET'):
        return {'status':'blocked','reason':'missing_remote_target','example':str(project/'AirPlan'/'state'/'airxdb'/'remote-device.env.example')}
    if not shutil.which('ssh'): return {'status':'blocked','reason':'ssh_not_found'}
    probe = run(env, 'printf ok', 20)
    if probe.returncode: return {'status':'blocked','reason':'ssh_probe_failed','stderr':probe.stderr.strip()}
    wd = workdir(env); mk = run(env, f'mkdir -p {q(wd)}', 20)
    if mk.returncode: return {'status':'blocked','reason':'remote_workdir_failed','stderr':mk.stderr.strip()}
    tool = detect_tool(env); install = 'skipped'
    if not tool and auto:
        inst = install_tool(env); install = 'ok' if inst.returncode == 0 else f'failed:{inst.returncode}'; tool = detect_tool(env)
    upd = {'AIRXDB_REMOTE_WORKDIR': wd}
    if tool: upd['AIRXDB_REMOTE_SCREENSHOT_TOOL'] = tool
    write_env(project/'AirPlan'/'state'/'airxdb'/'remote-device.env', upd)
    return {'status':'ok' if tool else 'blocked','target':env.get('AIRXDB_REMOTE_SSH_TARGET',''),'workdir':wd,'screenshotTool':tool,'autoConfigure':install,'hint':'' if tool else 'Install scrot/gnome-screenshot/grim/spectacle/import/screencapture or set AIRXDB_REMOTE_SCREENSHOT_TOOL.'}

def shot_cmd(tool, out, display):
    pre = f'export DISPLAY={q(display)}; ' if display else ''
    table = {'gnome-screenshot':f'gnome-screenshot -f {q(out)}','spectacle':f'spectacle -b -n -o {q(out)}','scrot':f'scrot {q(out)}','grim':f'grim {q(out)}','import':f'import -window root {q(out)}','screencapture':f'screencapture -x {q(out)}'}
    if tool not in table: raise SystemExit(f'unsupported screenshot tool: {tool}')
    return (pre if tool != 'screencapture' else '') + table[tool]

def emit(prefix, data):
    print(f'{prefix}_status={data.get("status","unknown")}')
    for k, v in data.items():
        if k != 'status': print(f'{k}={json.dumps(v, ensure_ascii=False) if isinstance(v,(dict,list)) else v}')

def main():
    ap = argparse.ArgumentParser(); ap.add_argument('--project', default='.'); ap.add_argument('--action', choices=['setup','status','screenshot'], default='setup'); ap.add_argument('--output-dir', default=''); ap.add_argument('--timeout', type=int, default=60); ap.add_argument('--no-auto-configure', action='store_true'); a = ap.parse_args()
    project = Path(a.project).expanduser().resolve(); ensure_files(project); env = read_env(project/'AirPlan'/'state'/'airxdb'/'remote-device.env')
    info = setup(project, env, auto=(a.action!='status' and not a.no_auto_configure))
    if a.action in ['setup','status'] or info.get('status') != 'ok': emit('airxdb_remote', info); raise SystemExit(0 if info.get('status')=='ok' else 2)
    outdir = Path(a.output_dir).expanduser().resolve() if a.output_dir else project/'AirPlan'/'docs'/'debug'/'airxdb-artifacts'; outdir.mkdir(parents=True, exist_ok=True)
    remote = info['workdir'].rstrip('/') + '/' + stamp() + '-airxdb-remote.png'; cap = run(env, shot_cmd(info['screenshotTool'], remote, env.get('AIRXDB_REMOTE_DISPLAY','')), a.timeout)
    data = {'status':'ok' if cap.returncode == 0 else 'failed','target':info['target'],'tool':info['screenshotTool'],'remoteFile':remote,'stdout':cap.stdout.strip(),'stderr':cap.stderr.strip(),'capturedAt':iso()}
    if cap.returncode == 0:
        b64 = run(env, f'base64 < {q(remote)}', a.timeout)
        local = outdir/(stamp()+'-airxdb-remote.png'); local.write_bytes(base64.b64decode(''.join(b64.stdout.split()))); data['screenshot'] = str(local)
    report = outdir/(stamp()+'-airxdb-remote-device.json'); report.write_text(json.dumps(data, ensure_ascii=False, indent=2)+'\n', encoding='utf-8'); data['report'] = str(report)
    if data['status'] == 'ok':
        log = project/'AirPlan'/'docs'/'debug'/'gui-debug-log.md'; log.parent.mkdir(parents=True, exist_ok=True); log.open('a', encoding='utf-8', newline='\n').write(f"\n## {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}: AirXDB remote screenshot\n\n- Remote target: `{data['target']}`\n- Screenshot: `{data.get('screenshot','')}`\n- Report: `{data['report']}`\n- AirDbg handoff: TODO\n- Residual risk: remote screenshots may contain sensitive data.\n")
    emit('airxdb_remote', data); raise SystemExit(0 if data['status']=='ok' else 1)
if __name__ == '__main__': main()