brain.py

import os
import json
import threading
import time
import requests
from datetime import datetime
from flask import Flask, jsonify, request, render_template, render_template_string

app = Flask(__name__)
PORT = int(os.environ.get("PORT", 7860))

WORKER_URL = os.environ.get("WORKER_URL", "")
GH_TOKEN_A = os.environ.get("GH_TOKEN_A", "")
GH_TOKEN_B = os.environ.get("GH_TOKEN_B", "")
GH_OWNER_A = os.environ.get("GH_OWNER_A", "")
GH_OWNER_B = os.environ.get("GH_OWNER_B", "")

# Mutable state (in-process; Brain is single-instance)
state = {
    "active": "A",
    "switched_at": datetime.utcnow().isoformat(),
    "desktop_a_url": os.environ.get("DESKTOP_A_URL", ""),
    "desktop_b_url": os.environ.get("DESKTOP_B_URL", ""),
}


# ── GitHub helpers ────────────────────────────────────────────────────────────

def _gh(token):
    return {
        "Authorization": f"Bearer {token}",
        "Accept": "application/vnd.github+json",
        "X-GitHub-Api-Version": "2022-11-28",
    }


def gh_codespaces(token):
    try:
        r = requests.get("https://api.github.com/user/codespaces",
                         headers=_gh(token), timeout=10)
        if r.status_code == 200:
            return r.json().get("codespaces", [])
    except Exception:
        pass
    return []


def gh_billing(token, owner):
    try:
        r = requests.get(
            f"https://api.github.com/users/{owner}/settings/billing/codespaces",
            headers=_gh(token), timeout=10,
        )
        if r.status_code == 200:
            return r.json()
    except Exception:
        pass
    return {}


def gh_codespace_port_url(token, port):
    """Discover the public forwarded URL for a port on the running codespace."""
    for cs in gh_codespaces(token):
        if cs.get("state") != "Available":
            continue
        name = cs["name"]
        try:
            r = requests.get(
                f"https://api.github.com/user/codespaces/{name}/ports",
                headers=_gh(token), timeout=10,
            )
            if r.status_code == 200:
                for p in r.json().get("ports", []):
                    if p.get("port_number") == port and p.get("visibility") == "public":
                        return p.get("browser_url", "").rstrip("/")
        except Exception:
            pass
        # Fallback: construct URL from codespace name
        return f"https://{name}-{port}.preview.app.github.dev"
    return ""


def gh_start(token, cs_name):
    try:
        r = requests.post(
            f"https://api.github.com/user/codespaces/{cs_name}/start",
            headers=_gh(token), timeout=15,
        )
        return r.status_code in (200, 202)
    except Exception:
        return False


def gh_stop(token, cs_name):
    try:
        r = requests.post(
            f"https://api.github.com/user/codespaces/{cs_name}/stop",
            headers=_gh(token), timeout=15,
        )
        return r.status_code in (200, 202)
    except Exception:
        return False


def hours_left(billing):
    included = billing.get("included_minutes", 120)
    used = billing.get("total_minutes_used", 0)
    return max(0, round((included - used) / 2, 1))  # 2-core → divide by 2


# ── Service health ────────────────────────────────────────────────────────────

def check(url, path="/api/health"):
    if not url:
        return "unconfigured"
    try:
        r = requests.get(url + path, timeout=5)
        return "online" if r.status_code == 200 else "error"
    except Exception:
        return "offline"


def active_desktop_url():
    key = "desktop_a_url" if state["active"] == "A" else "desktop_b_url"
    url = state.get(key, "")
    # Auto-discover if empty
    if not url:
        token = GH_TOKEN_A if state["active"] == "A" else GH_TOKEN_B
        if token:
            url = gh_codespace_port_url(token, 8080)
            if url:
                state[key] = url
    return url


# ── Keep-alive ────────────────────────────────────────────────────────────────

def _keep_alive():
    while True:
        time.sleep(600)
        for url in [state["desktop_a_url"], state["desktop_b_url"], WORKER_URL]:
            if url:
                try:
                    requests.get(url + "/api/health", timeout=5)
                except Exception:
                    pass

threading.Thread(target=_keep_alive, daemon=True).start()


# ── Routes ────────────────────────────────────────────────────────────────────

@app.route("/health")
def health():
    return jsonify({"status": "ok", "service": "brain", "active": state["active"]})


@app.route("/")
def dashboard():
    import os as _os
    tmpl_path = _os.path.join(_os.path.dirname(__file__), "templates", "index.html")
    with open(tmpl_path) as f:
        return render_template_string(f.read())


@app.route("/api/status")
def api_status():
    token_a = GH_TOKEN_A
    token_b = GH_TOKEN_B

    # Discover URLs from GitHub API if not set
    url_a = state["desktop_a_url"] or (gh_codespace_port_url(token_a, 8080) if token_a else "")
    url_b = state["desktop_b_url"] or (gh_codespace_port_url(token_b, 8080) if token_b else "")
    if url_a:
        state["desktop_a_url"] = url_a
    if url_b:
        state["desktop_b_url"] = url_b

    novnc_a = url_a.replace("-8080.", "-6080.") if url_a else ""
    novnc_b = url_b.replace("-8080.", "-6080.") if url_b else ""

    cs_a = gh_codespaces(token_a) if token_a else []
    cs_b = gh_codespaces(token_b) if token_b else []
    bill_a = gh_billing(token_a, GH_OWNER_A) if token_a and GH_OWNER_A else {}
    bill_b = gh_billing(token_b, GH_OWNER_B) if token_b and GH_OWNER_B else {}

    return jsonify({
        "active": state["active"],
        "switched_at": state["switched_at"],
        "desktop_a": {
            "api_url": url_a,
            "novnc_url": novnc_a,
            "health": check(url_a),
            "codespaces": [{"name": c["name"], "state": c["state"]} for c in cs_a],
            "hours_left": hours_left(bill_a),
        },
        "desktop_b": {
            "api_url": url_b,
            "novnc_url": novnc_b,
            "health": check(url_b),
            "codespaces": [{"name": c["name"], "state": c["state"]} for c in cs_b],
            "hours_left": hours_left(bill_b),
        },
        "worker": {"url": WORKER_URL, "health": check(WORKER_URL, "/health")},
    })


@app.route("/api/switch", methods=["POST"])
def api_switch():
    data = request.json or {}
    target = data.get("to")
    state["active"] = target if target in ("A", "B") else ("B" if state["active"] == "A" else "A")
    state["switched_at"] = datetime.utcnow().isoformat()
    return jsonify({"active": state["active"], "switched_at": state["switched_at"]})


@app.route("/api/codespace/start", methods=["POST"])
def api_cs_start():
    data = request.json or {}
    account = data.get("account", state["active"])
    token = GH_TOKEN_A if account == "A" else GH_TOKEN_B
    for cs in gh_codespaces(token):
        if cs["state"] in ("Shutdown", "Stopped", "ShuttingDown"):
            ok = gh_start(token, cs["name"])
            return jsonify({"ok": ok, "codespace": cs["name"]})
    return jsonify({"ok": False, "error": "no stopped codespace found"}), 404


@app.route("/api/codespace/stop", methods=["POST"])
def api_cs_stop():
    data = request.json or {}
    account = data.get("account", state["active"])
    token = GH_TOKEN_A if account == "A" else GH_TOKEN_B
    for cs in gh_codespaces(token):
        if cs["state"] == "Available":
            ok = gh_stop(token, cs["name"])
            return jsonify({"ok": ok, "codespace": cs["name"]})
    return jsonify({"ok": False, "error": "no running codespace found"}), 404


@app.route("/api/set-url", methods=["POST"])
def api_set_url():
    """Manually set a desktop URL (called by App after Codespace starts)."""
    data = request.json or {}
    account = data.get("account", "A")
    url = data.get("url", "").rstrip("/")
    key = "desktop_a_url" if account == "A" else "desktop_b_url"
    state[key] = url
    return jsonify({"ok": True, "account": account, "url": url})


@app.route("/proxy/<path:subpath>", methods=["GET", "POST"])
def proxy(subpath):
    url = active_desktop_url()
    if not url:
        return jsonify({"error": "no active desktop"}), 503
    target = f"{url}/{subpath}"
    try:
        if request.method == "POST":
            r = requests.post(target, json=request.get_json(), timeout=30)
        else:
            r = requests.get(target, timeout=30)
        return r.content, r.status_code, {"Content-Type": r.headers.get("Content-Type", "application/json")}
    except Exception as e:
        return jsonify({"error": str(e)}), 502


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=PORT, debug=False)