#!/usr/bin/env python3

# Miracle Claw (miracle_claw)
# Copyright (c) 2026 Milagro Distribution Corp.
# All rights reserved. See LICENSE file for terms.
"""mlg-serve — expose miracle_claw in a browser.

Usage:
    mlg-serve                       Start on default port 7777
    mlg-serve --port 8888           Different port
    mlg-serve --host 0.0.0.0        Bind all interfaces (LAN access)
    mlg-serve --no-auth             Disable token check (local-only)

Then open http://localhost:7777 in a browser.
"""
from __future__ import annotations
import argparse
import asyncio
import json
import os
import pty
import secrets
import signal
import sys
from pathlib import Path

try:
    from aiohttp import web
except ImportError:
    print("❌  Need 'aiohttp' package. Install: pip install aiohttp", file=sys.stderr)
    sys.exit(1)


HERE = Path(__file__).resolve().parent
DEFAULT_PORT = 7777
AUTH_TOKEN = ""
SESSIONS: dict[str, dict] = {}  # session_id -> {started, remote, cwd, bytes_in, bytes_out, pid}
MLG_SCRIPT = HERE / "mlg"


HTML_TEMPLATE = """<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>⚡ miracle_claw</title>
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/xterm@5.3.0/css/xterm.css">
  <style>
    :root {
      --bg: #0b0d12;
      --fg: #d8dee9;
      --accent: #00d9ff;
      --dim: #4c566a;
    }
    * { box-sizing: border-box; }
    html, body {
      margin: 0; padding: 0; height: 100%;
      background: var(--bg); color: var(--fg);
      font-family: ui-monospace, 'SF Mono', Menlo, Consolas, monospace;
    }
    #topbar {
      display: flex; align-items: center; justify-content: space-between;
      padding: 8px 14px; background: #1a1d24;
      border-bottom: 1px solid #2e3440; font-size: 12px;
    }
    #topbar .title { color: var(--accent); font-weight: 600; }
    #topbar .status { color: var(--dim); }
    #topbar .actions button {
      background: transparent; border: 1px solid var(--dim);
      color: var(--fg); padding: 3px 9px; border-radius: 4px;
      font-family: inherit; font-size: 11px; cursor: pointer; margin-left: 6px;
    }
    #topbar .actions button:hover { border-color: var(--accent); color: var(--accent); }
    #term { position: absolute; top: 36px; left: 0; right: 0; bottom: 0; padding: 8px; }
  </style>
</head>
<body>
  <div id="auth-banner" style="display:none;background:#f1fa8c;color:#2e343a;padding:6px 12px;text-align:center;font-size:12px;">
      Auth required. Open this page with <code>?token=&lt;TOKEN&gt;</code> appended to the URL.
    </div>
    <div id="topbar">
    <div>
      <span class="title">⚡ miracle_claw</span>
      <span class="status" id="status">  connected</span>
    </div>
    <div class="actions">
      <select id="theme-select" onchange="applyTheme(this.value)" style="background:transparent;color:inherit;border:1px solid var(--dim);border-radius:4px;font-family:inherit;font-size:11px;padding:2px 6px;">
        <option value="nord">nord</option>
        <option value="solarized">solarized</option>
        <option value="light">light</option>
        <option value="dracula">dracula</option>
      </select>
      <button onclick="term.clear()">clear</button>
      <button onclick="term.write('\\x03')">^C</button>
      <button onclick="location.reload()">reconnect</button>
    </div>
  </div>
  <div id="term"></div>

  <script src="https://cdn.jsdelivr.net/npm/xterm@5.3.0/lib/xterm.js"></script>
  <script src="https://cdn.jsdelivr.net/npm/xterm-addon-fit@0.8.0/lib/xterm-addon-fit.js"></script>
  <script>
    const TOKEN = '__TOKEN__';
    // Prefer the token from the URL if present (allows sharing links without restarts)
    const URL_TOKEN = new URLSearchParams(location.search).get('token') || '';
    const EFFECTIVE_TOKEN = URL_TOKEN || TOKEN;
    let term = null, fitAddon = null, ws = null;

    const THEMES = {
      nord: {
        background: '#0b0d12', foreground: '#d8dee9', cursor: '#00d9ff',
        black: '#3b4252', red: '#bf616a', green: '#a3be8c',
        yellow: '#ebcb8b', blue: '#81a1c1', magenta: '#b48ead',
        cyan: '#88c0d0', white: '#e5e9f0',
        brightBlack: '#4c566a', brightRed: '#bf616a', brightGreen: '#a3be8c',
        brightYellow: '#ebcb8b', brightBlue: '#81a1c1', brightMagenta: '#b48ead',
        brightCyan: '#8fbcbb', brightWhite: '#eceff4',
        bodyBg: '#0b0d12', topBg: '#1a1d24', accent: '#00d9ff', dim: '#4c566a',
      },
      solarized: {
        background: '#002b36', foreground: '#839496', cursor: '#93a1a1',
        black: '#073642', red: '#dc322f', green: '#859900',
        yellow: '#b58900', blue: '#268bd2', magenta: '#d33682',
        cyan: '#2aa198', white: '#eee8d5',
        brightBlack: '#586e75', brightRed: '#cb4b16', brightGreen: '#586e75',
        brightYellow: '#657b83', brightBlue: '#839496', brightMagenta: '#6c71c4',
        brightCyan: '#93a1a1', brightWhite: '#fdf6e3',
        bodyBg: '#002b36', topBg: '#073642', accent: '#b58900', dim: '#586e75',
      },
      light: {
        background: '#fafafa', foreground: '#1a1a1a', cursor: '#0066cc',
        black: '#1a1a1a', red: '#c91b1b', green: '#00a854',
        yellow: '#d97706', blue: '#0066cc', magenta: '#9333ea',
        cyan: '#0891b2', white: '#fafafa',
        brightBlack: '#666666', brightRed: '#ef4444', brightGreen: '#10b981',
        brightYellow: '#f59e0b', brightBlue: '#3b82f6', brightMagenta: '#a855f7',
        brightCyan: '#06b6d4', brightWhite: '#ffffff',
        bodyBg: '#fafafa', topBg: '#f0f0f0', accent: '#0066cc', dim: '#888888',
      },
      dracula: {
        background: '#282a36', foreground: '#f8f8f2', cursor: '#f8f8f0',
        black: '#21222c', red: '#ff5555', green: '#50fa7b',
        yellow: '#f1fa8c', blue: '#bd93f9', magenta: '#ff79c6',
        cyan: '#8be9fd', white: '#f8f8f2',
        brightBlack: '#6272a4', brightRed: '#ff6e6e', brightGreen: '#69ff94',
        brightYellow: '#ffffa5', brightBlue: '#d6acff', brightMagenta: '#ff92df',
        brightCyan: '#a4ffff', brightWhite: '#ffffff',
        bodyBg: '#282a36', topBg: '#44475a', accent: '#bd93f9', dim: '#6272a4',
      },
    };

    function applyTheme(name) {
      const t = THEMES[name] || THEMES.nord;
      document.body.style.background = t.bodyBg;
      document.body.style.color = t.foreground;
      document.getElementById('topbar').style.background = t.topBg;
      document.querySelectorAll('#topbar .title').forEach(el => el.style.color = t.accent);
      if (term) {
        term.options.theme = {
          background: t.background, foreground: t.foreground, cursor: t.cursor,
          black: t.black, red: t.red, green: t.green, yellow: t.yellow,
          blue: t.blue, magenta: t.magenta, cyan: t.cyan, white: t.white,
          brightBlack: t.brightBlack, brightRed: t.brightRed,
          brightGreen: t.brightGreen, brightYellow: t.brightYellow,
          brightBlue: t.brightBlue, brightMagenta: t.brightMagenta,
          brightCyan: t.brightCyan, brightWhite: t.brightWhite,
        };
      }
      localStorage.setItem('miracle_claw_theme', name);
    }

    const savedTheme = localStorage.getItem('miracle_claw_theme') || 'nord';

    term = new Terminal({
      cursorBlink: true, fontSize: 13,
      fontFamily: 'ui-monospace, "SF Mono", Menlo, Consolas, monospace',
      theme: THEMES.nord,  // overridden by applyTheme()
      scrollback: 5000, convertEol: true,
    });
    fitAddon = new FitAddon.FitAddon();
    term.loadAddon(fitAddon);
    term.open(document.getElementById('term'));
    fitAddon.fit();
    const AUTH_REQUIRED = '__AUTH_REQUIRED__' === 'true';
    if (AUTH_REQUIRED && !EFFECTIVE_TOKEN) {
      document.getElementById('auth-banner').style.display = 'block';
    }
    applyTheme(savedTheme);
    document.getElementById('theme-select').value = savedTheme;
    window.term = term;

    const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
    const wsUrl = proto + '//' + location.host + '/ws?token=' + encodeURIComponent(EFFECTIVE_TOKEN);
    ws = new WebSocket(wsUrl);

    ws.onopen = () => {
      const cols = term.cols, rows = term.rows;
      ws.send(JSON.stringify({type: 'resize', cols, rows}));
    };
    ws.onmessage = (e) => {
      term.write(e.data);
    };
    ws.onclose = () => {
      document.getElementById('status').textContent = '  disconnected';
      term.write('\\r\\n\\x1b[33m[connection closed — refresh to reconnect]\\x1b[0m\\r\\n');
    };
    ws.onerror = (e) => { console.error(e); };

    term.onData((data) => {
      if (ws && ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({type: 'input', data}));
      }
    });

    window.addEventListener('resize', () => {
      fitAddon.fit();
      if (ws && ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({type: 'resize', cols: term.cols, rows: term.rows}));
      }
    });
  </script>
</body>
</html>
"""


# ─── PTY plumbing ───────────────────────────────────────────────────────────
def set_winsize(fd: int, rows: int, cols: int):
    import fcntl, struct, termios
    size = struct.pack("HHHH", rows, cols, 0, 0)
    fcntl.ioctl(fd, termios.TIOCSWINSZ, size)


def set_nonblocking(fd: int) -> None:
    """Make a fd non-blocking. NOT USED — kept for reference.
    PTY masters work better in blocking mode with select() than non-blocking mode."""
    import fcntl
    flags = fcntl.fcntl(fd, fcntl.F_GETFL)
    fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)


async def relay_pty(master: int, ws: web.WebSocketResponse, rec_file=None, sid: str = "") -> None:
    """Forward bytes from PTY master → WebSocket. Optionally record to asciinema .cast file."""
    import selectors
    loop = asyncio.get_running_loop()
    # Reuse the time origin set by ws_handler so input and output events
    # in the recording share the same t=0.
    t0 = getattr(ws, "_mlg_t0", loop.time())
    sel = selectors.DefaultSelector()
    sel.register(master, selectors.EVENT_READ)
    while True:
        # Run a blocking select on a thread; aiohttp can do its other work meanwhile.
        # Returns list of (key, mask) ready events, or [] after timeout.
        ready = await loop.run_in_executor(None, sel.select, 0.1)
        if not ready:
            await asyncio.sleep(0.01)
            continue
        try:
            data = os.read(master, 4096)
        except OSError as e:
            print(f"[pty] OSError on read: {e}", flush=True)
            break
        if not data:
            print("[pty] EOF", flush=True)
            break
        if ws.closed:
            print("[pty] ws closed, stop relay", flush=True)
            break
        text = data.decode("utf-8", errors="replace")
        print(f"[relay] {len(data)}b: {text[:80]!r}", flush=True)
        if rec_file:
            elapsed = loop.time() - t0
            try:
                import json as _json
                rec_file.write(_json.dumps([round(elapsed, 6), "o", text]) + "\n")
                rec_file.flush()
            except Exception:
                pass
        try:
            await ws.send_str(text)
            if sid and sid in SESSIONS:
                SESSIONS[sid]["bytes_out"] += len(data)
        except Exception as e:
            print(f"[relay] ws.send failed ({type(e).__name__}: {e}); closing relay", flush=True)
            try:
                await ws.close(code=1011, message=b"relay send failed")
            except Exception:
                pass
            return


async def ws_handler(request: web.Request) -> web.WebSocketResponse:
    """One WebSocket connection: spawn mlg in a PTY, bridge it."""
    # Auth check
    if AUTH_TOKEN:
        token = request.query.get("token", "")
        if token != AUTH_TOKEN:
            return web.json_response(
                {"error": "bad token", "hint": "open this page with ?token=<token> appended to the URL"},
                status=401,
            )

    ws = web.WebSocketResponse(max_msg_size=2**20)
    await ws.prepare(request)

    # Spawn shell with mlg in a PTY first (we need pid for SESSIONS dict)
    pid, master = pty.fork()
    if pid == 0:
        # Child process
        os.environ["TERM"] = "xterm-256color"
        os.environ["COLUMNS"] = "120"
        os.environ["LINES"] = "40"
        # Sentinel so mlg knows it's being driven by a browser PTY (not a
        # real terminal). mlg uses this to *skip* the alt-screen handoff
        # — xterm.js renders the whole frame every keystroke, so 1049
        # would just produce visible garbage when the buffer flips.
        os.environ["MILAGRO_CHAT_SURFACE"] = "mlg-serve"
        # Ensure MILAGRO_API_KEY is inherited (in case parent lacks it)
        if "MILAGRO_API_KEY" not in os.environ:
            secrets_path = Path.home() / ".config" / "secrets" / "milagro.env"
            if secrets_path.exists():
                for line in secrets_path.read_text().splitlines():
                    if line.startswith("export "):
                        k, _, v = line[7:].partition("=")
                        os.environ[k] = v
        try:
            os.chdir(os.path.expanduser("~"))
        except OSError:
            pass
        # Exec mlg in this PTY. We use the script's own python (-u for unbuffered),
        # passing the mlg script path explicitly so it runs no matter the cwd.
        import shutil as _shutil
        py = _shutil.which("python3") or "/usr/bin/python3"
        # Browser sessions default to --yes so SAFE + NORMAL bash commands
        # run without prompting. DANGEROUS still requires YES-EXEC.
        # Set MILAGRO_NO_AUTO=1 in the environment to opt out.
        if os.environ.get("MILAGRO_NO_AUTO"):
            os.execv(py, [py, "-u", str(MLG_SCRIPT)])
        else:
            os.execv(py, [py, "-u", str(MLG_SCRIPT), "--yes"])

    # Parent: set initial size, register session, start relay
    set_winsize(master, 40, 120)

    import datetime as dt, uuid as _uuid
    sid = _uuid.uuid4().hex
    # Recording time origin — used by both the relay and the input-event
    # recorder so timestamps line up. We use the asyncio event loop's clock.
    t0 = asyncio.get_event_loop().time()
    ws._mlg_t0 = t0  # share with relay_pty
    SESSIONS[sid] = {
        "started": dt.datetime.now(dt.timezone.utc).isoformat(),
        "remote": f"{request.remote}",
        "cwd": "~",
        "bytes_in": 0,
        "bytes_out": 0,
        "pid": pid,
    }

    # Optional recording (always on for now; add --no-record flag later)
    rec_file = None
    if True:  # toggle this if you want to disable
        recordings_dir = HERE / "recordings"
        recordings_dir.mkdir(exist_ok=True)
        import datetime as dt
        rec_name = dt.datetime.now().strftime("%Y%m%d-%H%M%S") + "-" + sid[:8] + ".cast"
        rec_path = recordings_dir / rec_name
        rec_file = open(rec_path, "w", encoding="utf-8")
        # asciinema v2 header
        import json as _json
        rec_file.write(_json.dumps({
            "version": 2,
            "width": 120, "height": 40,
            "timestamp": int(dt.datetime.now().timestamp()),
            "title": f"miracle_claw session {sid[:8]}",
            "env": {"TERM": "xterm-256color", "SHELL": "/bin/bash"},
        }) + "\n")
        rec_file.flush()
        if sid in SESSIONS:
            SESSIONS[sid]["recording"] = rec_name

    relay_task = asyncio.create_task(relay_pty(master, ws, rec_file, sid))
    # Log relay task failures instead of silently killing the WS session.
    # (Previously a NameError in relay_pty closed the relay but the WS stayed
    # open showing only what had already been flushed — see MEMORY 2026-06-27.)
    def _relay_done(t: asyncio.Task) -> None:
        if t.cancelled():
            return
        exc = t.exception()
        if exc is not None:
            print(f"[relay-task] CRASHED: {type(exc).__name__}: {exc}", flush=True)
    relay_task.add_done_callback(_relay_done)

    try:
        async for msg in ws:
            if msg.type == web.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                except json.JSONDecodeError:
                    print(f"[ws] non-JSON msg: {msg.data[:60]!r}", flush=True)
                    continue
                if data.get("type") == "input":
                    try:
                        payload = data.get("data", "").encode("utf-8")
                        print(f"[ws→pty] {payload!r}", flush=True)
                        if rec_file:
                            # Record input events too (asciinema v2 supports
                            # 'i' = input from user, 'o' = output to user).
                            # Without this we can't replay keystrokes.
                            import json as _json_input
                            elapsed_input = asyncio.get_event_loop().time() - t0
                            try:
                                rec_file.write(_json_input.dumps([round(elapsed_input, 6), "i", payload.decode("utf-8", errors="replace")]) + "\n")
                                rec_file.flush()
                            except Exception:
                                pass
                        os.write(master, payload)
                        if sid in SESSIONS:
                            SESSIONS[sid]["bytes_in"] += len(payload)
                    except OSError as e:
                        print(f"[ws] OSError on write: {e}", flush=True)
                        break
                elif data.get("type") == "resize":
                    print(f"[serve] resize to {data.get('rows')}x{data.get('cols')}, pid={pid}", flush=True)
                    set_winsize(master, data.get("rows", 40), data.get("cols", 120))
                    # Notify child so it can re-layout its UI (the input box
                    # position depends on terminal height).
                    try:
                        import signal as _sig
                        os.kill(pid, _sig.SIGWINCH)
                        print(f"[serve] SIGWINCH sent to pid={pid}", flush=True)
                    except OSError as e:
                        print(f"[serve] SIGWINCH failed: {e}", flush=True)
            elif msg.type == web.WSMsgType.ERROR:
                break
    except Exception:
        pass
    finally:
        print(f"[ws] session {sid[:8]} ending (ws.closed={ws.closed}, code={ws.close_code if hasattr(ws, 'close_code') else '?'})", flush=True)
        relay_task.cancel()
        SESSIONS.pop(sid, None)
        try:
            if rec_file:
                rec_file.close()
        except Exception:
            pass
        try:
            os.kill(pid, signal.SIGTERM)
        except (ProcessLookupError, PermissionError):
            pass
        try:
            os.close(master)
        except OSError:
            pass
        try:
            wpid, status = os.waitpid(pid, os.WNOHANG)
            if wpid:
                print(f"[ws] child {wpid} reaped, status={status}", flush=True)
        except ChildProcessError:
            pass

    return ws


async def index_handler(request: web.Request) -> web.Response:
    # If a ?token= is supplied in URL, strip it (we don't want it embedded in HTML
    # for the WS connection to reuse - the URL is the source of truth)
    import copy
    body = HTML_TEMPLATE.replace("__TOKEN__", AUTH_TOKEN or "")
    # Tell the client whether auth is required so it can show a banner
    auth_required = "true" if AUTH_TOKEN else "false"
    body = body.replace("__AUTH_REQUIRED__", auth_required)
    return web.Response(text=body, content_type="text/html")


async def healthz_handler(request: web.Request) -> web.Response:
    return web.json_response({"status": "ok", "sessions": len(SESSIONS)})


async def sessions_handler(request: web.Request) -> web.Response:
    """List active sessions (admin endpoint)."""
    if AUTH_TOKEN:
        token = request.query.get("token", "")
        if token != AUTH_TOKEN:
            return web.Response(status=401, text="bad token")
    out = []
    for sid, s in SESSIONS.items():
        out.append({
            "id": sid[:12],
            "started": s["started"],
            "remote": s["remote"],
            "cwd": s.get("cwd", "?"),
            "bytes_in": s["bytes_in"],
            "bytes_out": s["bytes_out"],
        })
    return web.json_response({"sessions": out, "count": len(out)})


async def record_handler(request: web.Request) -> web.Response:
    """List recorded sessions available for replay."""
    if AUTH_TOKEN:
        token = request.query.get("token", "")
        if token != AUTH_TOKEN:
            return web.Response(status=401, text="bad token")
    recordings_dir = HERE / "recordings"
    if not recordings_dir.exists():
        return web.json_response({"recordings": []})
    recs = []
    for f in sorted(recordings_dir.glob("*.cast"), reverse=True)[:50]:
        recs.append({
            "filename": f.name,
            "size_bytes": f.stat().st_size,
            "modified": f.stat().st_mtime,
        })
    return web.json_response({"recordings": recs})


def main():
    global AUTH_TOKEN

    p = argparse.ArgumentParser(
        prog="mlg-serve",
        description="Serve miracle_claw in a browser via WebSocket + xterm.js.",
    )
    p.add_argument("--host", default="127.0.0.1", help="Bind address (default: localhost only)")
    p.add_argument("--port", type=int, default=DEFAULT_PORT, help=f"Port (default: {DEFAULT_PORT})")
    p.add_argument("--no-auth", action="store_true", help="Disable token auth (local-only)")
    p.add_argument("--token", help="Custom auth token (default: random)")
    args = p.parse_args()

    if args.no_auth:
        AUTH_TOKEN = ""
    elif args.token:
        AUTH_TOKEN = args.token
    else:
        AUTH_TOKEN = secrets.token_urlsafe(16)

    app = web.Application()
    app.router.add_get("/", index_handler)
    app.router.add_get("/healthz", healthz_handler)
    app.router.add_get("/sessions", sessions_handler)
    app.router.add_get("/recordings", record_handler)
    app.router.add_get("/ws", ws_handler)

    if AUTH_TOKEN:
        print(f"🔒  Auth token: {AUTH_TOKEN}")
        print(f"    Open: http://{args.host}:{args.port}/?token={AUTH_TOKEN}")
    else:
        print(f"🔓  Auth disabled. Open: http://{args.host}:{args.port}/")
    print(f"⚡ Listening on http://{args.host}:{args.port}")

    web.run_app(app, host=args.host, port=args.port, print=lambda *_: None)


if __name__ == "__main__":
    main()