#!/usr/bin/env python3
"""
refresh-node-cookies.py: Extract fresh cookies from a running node's Chromium CDP
instance inside its isolated network namespace and save to ~/.config/muse-cli/<node>/cookies.txt.
"""
import sys
import os
import json
import urllib.request
from pathlib import Path

NODES = {
    "muse": 9410,
    "pip": 9420,
    "646": 9430,
    "opm": 9440,
    "def": 9450,
}

def export_cookies(node):
    port = NODES.get(node)
    if not port:
        print(f"Unknown node: {node}", file=sys.stderr)
        return False

    conf_dir = Path.home() / ".config" / "muse-cli" / node
    conf_dir.mkdir(parents=True, exist_ok=True)
    cookie_file = conf_dir / "cookies.txt"
    cfg_file = conf_dir / "config.json"

    # We import websocket inside netns execution
    try:
        import websocket
    except ImportError:
        print("websocket-client not installed", file=sys.stderr)
        return False

    try:
        req = urllib.request.urlopen(f"http://127.0.0.1:{port}/json", timeout=3)
        tabs = json.loads(req.read().decode())
        if not tabs:
            print(f"No tabs found on port {port}", file=sys.stderr)
            return False
        ws_url = tabs[0]["webSocketDebuggerUrl"]
        ws = websocket.create_connection(ws_url, timeout=5)
        ws.send(json.dumps({
            "id": 1,
            "method": "Network.getCookies",
            "params": {"urls": ["https://muse.ai"]}
        }))
        res = json.loads(ws.recv())
        ws.close()

        cookies = res.get("result", {}).get("cookies", [])
        if not cookies:
            print(f"No cookies returned from CDP for {node}", file=sys.stderr)
            return False

        lines = [f"{c['name']}={c['value']}" for c in cookies]
        cookie_file.write_text("; ".join(lines) + "\n")
        cookie_file.chmod(0o600)

        cfg_data = {}
        if cfg_file.exists():
            try:
                cfg_data = json.loads(cfg_file.read_text())
            except Exception:
                pass
        cfg_data["cookies_file"] = str(cookie_file)
        cfg_file.write_text(json.dumps(cfg_data, indent=2))
        cfg_file.chmod(0o600)
        return True
    except Exception as e:
        print(f"Failed to export cookies for {node}: {e}", file=sys.stderr)
        return False

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: refresh-node-cookies.py <node>", file=sys.stderr)
        sys.exit(1)
    success = export_cookies(sys.argv[1])
    sys.exit(0 if success else 1)