#!/usr/bin/env python3
"""rig-pocket: a phone-friendly status page and push notifier for OpenRig.

Reads state only through the documented `rig` CLI (`--json`), so it needs no
access to the daemon's internals. Standard library only.
"""
import glob
import json
import os
import re
import subprocess
import threading
import time
import urllib.parse
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

HERE = Path(__file__).resolve().parent
STATIC = HERE / "static"

RIG = os.environ.get("RIG_BIN", "rig")
BIND = os.environ.get("RIG_POCKET_BIND", "127.0.0.1")
PORT = int(os.environ.get("RIG_POCKET_PORT", "7440"))
POLL = int(os.environ.get("RIG_POCKET_POLL", "30"))
REPOS = os.environ.get("RIG_POCKET_REPOS", "")  # space separated globs of git work trees
ORIGIN = os.environ.get("RIG_POCKET_ORIGIN", "")  # e.g. https://rig.example.com, checked on POST
PUBLIC_URL = os.environ.get("RIG_POCKET_URL", "")  # where a notification tap should go
THEME_CSS = os.environ.get("RIG_POCKET_THEME_CSS", "https://assets.weircon.dk/css/weirsoe-theme.css")
THEME_JS = os.environ.get("RIG_POCKET_THEME_JS", "https://assets.weircon.dk/js/weirsoe-theme.js")
NTFY_URL = os.environ.get("NTFY_URL", "")
NTFY_TOPIC = os.environ.get("NTFY_TOPIC", "")
NTFY_TOKEN = os.environ.get("NTFY_TOKEN", "")
DONE_DEST = re.compile(os.environ.get("RIG_POCKET_DONE_DEST", r"^orch-lead@"))
HUMAN_DEST = re.compile(os.environ.get("RIG_POCKET_HUMAN_DEST", r"^operator-agent@kernel$"))

HUMAN_SEATS = re.compile(os.environ.get("RIG_POCKET_HUMAN_SEATS", r"^operator-human@"))
# the seat a decision is recorded as (rig derives the resolver from the seat environment)
HUMAN_ENV = {"OPENRIG_SESSION_NAME": os.environ.get("RIG_POCKET_HUMAN_SEAT", "operator-human@kernel")}

SEAT = re.compile(r"^[A-Za-z0-9._-]+@[A-Za-z0-9._-]+$")
QITEM = re.compile(r"^qitem-[A-Za-z0-9-]+$")
KEYS = {str(n) for n in range(1, 10)} | {"esc"}
CLOSED = {"done", "closed", "failed", "canceled", "cancelled"}

OPTION = re.compile(r"^\s*(❯\s*)?(\d)\.\s+(\S.*?)\s*$")
FRAME = re.compile(r"^[\s│┃|╌─━┄┈╭╮╰╯]+|[\s│┃|╌─━┄┈╭╮╰╯]+$")
BULLET = re.compile(r"^\s*●\s?")
INPUT_RULE = re.compile(r"^\s*─{8,}")


def clean(line):
    return FRAME.sub("", line).strip()


SPINNER = re.compile(r"\((?:(\d+)h )?(?:(\d+)m )?(\d+)s · [↓↑] ([\d.]+)(k?) tokens")
STALL_MIN = float(os.environ.get("RIG_POCKET_STALL_MIN", "15"))
STALL_RATE = float(os.environ.get("RIG_POCKET_STALL_RATE", "50"))  # tokens per minute
MEM_WARN = float(os.environ.get("RIG_POCKET_MEM_WARN", "0.9"))


def stalled(screen):
    """Minutes of a turn that has gone on long while barely producing tokens (a hung API stream), else None."""
    m = None
    for m in SPINNER.finditer(screen):
        pass
    if not m:
        return None
    h, mins, secs, tokens, k = m.groups()
    minutes = int(h or 0) * 60 + int(mins or 0) + int(secs) / 60
    tokens = float(tokens) * (1000 if k else 1)
    if minutes >= STALL_MIN and tokens / minutes < STALL_RATE:
        return round(minutes)
    return None


def memory():
    """Memory use of this container (cgroup v2), as a fraction of its limit, or None."""
    try:
        cur = int(Path("/sys/fs/cgroup/memory.current").read_text())
        limit = Path("/sys/fs/cgroup/memory.max").read_text().strip()
        if limit == "max":
            limit = next(int(l.split()[1]) * 1024 for l in open("/proc/meminfo") if l.startswith("MemTotal"))
        return {"used": cur, "limit": int(limit), "fraction": round(cur / int(limit), 3)}
    except (OSError, ValueError, StopIteration):
        return None


def parse_ask(screen, idle):
    """What a seat is waiting for, read off its screen.

    menu:     a numbered choice with the cursor (❯) on it (permission prompts, multiple choice questions)
    question: the agent stopped and its last message ends in a question
    """
    lines = [l for l in screen.splitlines() if not l.startswith("[rig capture:")]
    tail = lines[-30:]
    opts = [(i, OPTION.match(l)) for i, l in enumerate(tail)]
    opts = [(i, m) for i, m in opts if m]
    if opts and any(m.group(1) for _, m in opts):
        above = [clean(l) for l in tail[max(0, opts[0][0] - 8):opts[0][0]]]
        above = [t for t in above if t and not t.startswith(("Tip:", "Esc to"))]
        # the question line if there is one ("Do you want to proceed?"), else the nearest heading
        title = next((t for t in reversed(above) if t.endswith("?")), above[-1] if above else "")
        options = [{"key": m.group(2), "label": m.group(3)} for _, m in opts]
        return {"kind": "menu", "title": title, "options": options}

    if not idle:
        return None
    # the last agent message: from the last "●" line down to the input box rule
    end = next((i for i in range(len(lines) - 1, -1, -1) if INPUT_RULE.match(lines[i])), len(lines))
    end = next((i for i in range(end - 1, -1, -1) if INPUT_RULE.match(lines[i])), end)
    start = next((i for i in range(end - 1, -1, -1) if BULLET.match(lines[i])), None)
    if start is None:
        return None
    body = [lines[start]]
    for l in lines[start + 1:end]:
        if l and not l[0].isspace():  # status line (✻ …) or the input box: the message has ended
            break
        body.append(l)
    text = " ".join(BULLET.sub("", l).strip() for l in body if l.strip())
    text = re.sub(r"\s+", " ", text).strip()
    if text.endswith("?") and not text.startswith(("Bash(", "Read(", "Edit(", "Write(")):
        return {"kind": "question", "text": text[-600:]}
    return None


def rig(*args, timeout=20, env=None):
    try:
        out = subprocess.run([RIG, *args], capture_output=True, text=True, timeout=timeout,
                             env={**os.environ, **env} if env else None)
        return out.stdout
    except (OSError, subprocess.TimeoutExpired):
        return ""


def rig_json(*args):
    try:
        return json.loads(rig(*args) or "null")
    except json.JSONDecodeError:
        return None


def git(path, *args):
    try:
        return subprocess.run(["git", "-C", path, *args], capture_output=True, text=True, timeout=10).stdout
    except (OSError, subprocess.TimeoutExpired):
        return ""


def as_list(data):
    if isinstance(data, list):
        return data
    if isinstance(data, dict):
        return data.get("items") or data.get("entries") or []
    return []


def repo_paths():
    paths = []
    for pattern in REPOS.split():
        for p in sorted(glob.glob(os.path.expanduser(pattern))):
            if os.path.exists(os.path.join(p, ".git")):
                paths.append(p)
    return paths


def collect():
    rigs = []
    for r in as_list(rig_json("ps", "--json")):
        if r.get("isArchived"):
            continue
        seats = []
        for n in as_list(rig_json("ps", "--nodes", "--rig", r["name"], "--json")):
            name = n.get("canonicalSessionName", "")
            activity = (n.get("agentActivity") or {}).get("state", "unknown")
            # human seats are dashboards, not agents: nothing on them is a question for the person
            human = bool(HUMAN_SEATS.search(name))
            screen = rig("capture", name, "--lines", "40") if name and not human else ""
            ask = parse_ask(screen, activity == "idle") if screen else None
            seats.append({
                "name": name,
                "activity": activity,
                "status": n.get("sessionStatus"),
                "work": n.get("assignedWorkCount", 0),
                "lastActivity": n.get("lastActivity"),
                "waiting": ask is not None,
                "ask": ask,
                "stalled": stalled(screen) if screen and not ask else None,
            })
        rigs.append({"name": r["name"], "status": r.get("status"), "uptime": r.get("uptime"),
                     "attention": r.get("attentionCount", 0), "seats": seats})

    queue = []
    for q in as_list(rig_json("queue", "list", "-A", "-a", "--limit", "60", "--json")):
        dest = q.get("destinationSession") or ""
        queue.append({"id": q.get("qitemId"), "state": q.get("state"), "dest": dest,
                      "source": q.get("sourceSession"), "summary": q.get("summary"), "updated": q.get("tsUpdated"),
                      # a human gate: blocked on a human seat until someone records a decision (rig queue resolve)
                      "gate": q.get("state") == "blocked" and bool(HUMAN_SEATS.search(q.get("blockedOn") or ""))})

    repos = []
    for p in repo_paths():
        commits = []
        for line in git(p, "log", "--branches", "-15", "--format=%H%x09%h%x09%cI%x09%D%x09%s").splitlines():
            full, short, when, refs, subject = (line.split("\t") + [""] * 5)[:5]
            commits.append({"hash": full, "short": short, "when": when, "refs": refs, "subject": subject})
        repos.append({"name": os.path.basename(p), "commits": commits})

    return {"updated": time.time(), "rigs": rigs, "queue": queue, "repos": repos, "memory": memory()}


def notify(title, body, priority=3, tags=""):
    if not (NTFY_URL and NTFY_TOPIC):
        return
    headers = {"Title": title.encode("utf-8"), "Priority": str(priority)}
    if tags:
        headers["Tags"] = tags
    if PUBLIC_URL:
        headers["Click"] = PUBLIC_URL
    if NTFY_TOKEN:
        headers["Authorization"] = f"Bearer {NTFY_TOKEN}"
    req = urllib.request.Request(f"{NTFY_URL.rstrip('/')}/{NTFY_TOPIC}", data=body.encode("utf-8"), headers=headers)
    try:
        urllib.request.urlopen(req, timeout=10).close()
    except OSError:
        pass


class Watcher:
    """Turns successive snapshots into notifications. The first snapshot only seeds."""

    def __init__(self):
        self.seen_commits = None
        self.waiting_polls = {}
        self.notified_waiting = set()
        self.queue_state = None
        self.notified_gates = set()
        self.notified_stalled = set()
        self.mem_warned = False

    def events(self, snap):
        out = []
        commits = {c["hash"]: (r["name"], c) for r in snap["repos"] for c in r["commits"]}
        queue = {q["id"]: q for q in snap["queue"]}
        if self.seen_commits is None:
            self.seen_commits = set(commits)
            self.queue_state = {k: q["state"] for k, q in queue.items()}
        else:
            new = [commits[h] for h in commits if h not in self.seen_commits]
            self.seen_commits |= set(commits)
            by_repo = {}
            for repo, c in new:
                by_repo.setdefault(repo, []).append(c["subject"])
            for repo, subjects in by_repo.items():
                out.append((f"{repo}: {len(subjects)} ny commit" + ("s" if len(subjects) > 1 else ""),
                            "\n".join(subjects), 2, "package"))
            for k, q in queue.items():
                before = self.queue_state.get(k)
                if before is None and HUMAN_DEST.search(q["dest"] or "") and q["state"] not in CLOSED:
                    out.append(("Teamet spørger dig", q["summary"] or q["source"] or k, 4, "raising_hand"))
                if before and before not in CLOSED and q["state"] in CLOSED and DONE_DEST.search(q["dest"] or ""):
                    out.append((f"Opgave {q['state']}: {q['dest']}", q["summary"] or k, 4, "white_check_mark"))
            self.queue_state = {k: q["state"] for k, q in queue.items()}

            for k, q in queue.items():
                if q.get("gate") and k not in self.notified_gates:
                    out.append(("Beslutning til dig", q["summary"] or k, 4, "raising_hand"))
        self.notified_gates = {k for k, q in queue.items() if q.get("gate")}

        # a prompt must survive two polls before it pings, so prompts answered quickly stay quiet
        for r in snap["rigs"]:
            for s in r["seats"]:
                n = s["name"]
                self.waiting_polls[n] = self.waiting_polls.get(n, 0) + 1 if s["waiting"] else 0
                if self.waiting_polls[n] >= 2 and n not in self.notified_waiting:
                    self.notified_waiting.add(n)
                    out.append((f"{n} venter på svar", ask_text(s.get("ask")), 4, "warning"))
                if not s["waiting"]:
                    self.notified_waiting.discard(n)
                if s.get("stalled") and n not in self.notified_stalled:
                    self.notified_stalled.add(n)
                    out.append((f"{n} ser ud til at hænge",
                                f"Samme tur i {s['stalled']} min uden at producere noget. Åbn og send Esc.", 4, "hourglass"))
                if not s.get("stalled"):
                    self.notified_stalled.discard(n)

        mem = snap.get("memory")
        if mem and mem["fraction"] >= MEM_WARN and not self.mem_warned:
            self.mem_warned = True
            out.append(("Containeren er ved at løbe tør for hukommelse",
                        f"{mem['used'] / 2**30:.1f} af {mem['limit'] / 2**30:.1f} GB brugt. Agenterne kan gå i stå.", 5, "rotating_light"))
        elif mem and mem["fraction"] < MEM_WARN - 0.1:
            self.mem_warned = False
        return out


def ask_text(ask):
    if not ask:
        return "Åbn for at se skærmen og svare."
    if ask["kind"] == "menu":
        return (ask["title"] + "\n" + "\n".join(f"{o['key']}. {o['label']}" for o in ask["options"])).strip()
    return ask["text"]


STATE = {"snap": None}
LOCK = threading.Lock()


def poll_loop():
    watcher = Watcher()
    while True:
        try:
            snap = collect()
            with LOCK:
                STATE["snap"] = snap
            for title, body, prio, tags in watcher.events(snap):
                notify(title, body, prio, tags)
        except Exception as e:  # keep polling whatever one round did
            print("poll failed:", e, flush=True)
        time.sleep(POLL)


def known_seats():
    with LOCK:
        snap = STATE["snap"] or {"rigs": []}
    return {s["name"] for r in snap["rigs"] for s in r["seats"]}


def open_gates():
    with LOCK:
        snap = STATE["snap"] or {"queue": []}
    return {q["id"] for q in snap["queue"] if q.get("gate")}


TYPES = {".html": "text/html; charset=utf-8", ".js": "text/javascript; charset=utf-8",
         ".webmanifest": "application/manifest+json", ".svg": "image/svg+xml", ".png": "image/png",
         ".json": "application/json"}


def csp():
    ext = " ".join(sorted({u.split("/", 3)[0] + "//" + u.split("/", 3)[2] for u in (THEME_CSS, THEME_JS) if u.startswith("http")}))
    return (f"default-src 'self'; style-src 'self' {ext}; script-src 'self' {ext}; img-src 'self' data:; "
            "connect-src 'self'; manifest-src 'self'; worker-src 'self'; frame-ancestors 'none'; base-uri 'none'")


class Handler(BaseHTTPRequestHandler):
    server_version = "rig-pocket"

    def log_message(self, *args):
        pass

    def send(self, code, body, ctype="application/json", cache="no-store"):
        data = body if isinstance(body, bytes) else body.encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Cache-Control", cache)
        self.send_header("Content-Security-Policy", csp())
        self.send_header("X-Content-Type-Options", "nosniff")
        self.send_header("Referrer-Policy", "no-referrer")
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        path = self.path.split("?", 1)[0]
        if path == "/api/status":
            with LOCK:
                snap = STATE["snap"]
            return self.send(200, json.dumps(snap or {"updated": 0, "rigs": [], "queue": [], "repos": []}))
        if path.startswith("/api/capture/"):
            seat = urllib.parse.unquote(path.rsplit("/", 1)[1])
            if not SEAT.match(seat) or seat not in known_seats():
                return self.send(404, '{"error":"unknown seat"}')
            return self.send(200, json.dumps({"seat": seat, "screen": rig("capture", seat, "--lines", "60")}))
        name = "index.html" if path in ("/", "") else path.lstrip("/")
        f = (STATIC / name).resolve()
        if STATIC not in f.parents or not f.is_file():
            return self.send(404, "not found", "text/plain")
        data = f.read_bytes()
        if f.name == "index.html":
            data = data.replace(b"{{THEME_CSS}}", THEME_CSS.encode()).replace(b"{{THEME_JS}}", THEME_JS.encode())
        cache = "no-cache" if f.suffix in (".html", ".webmanifest") or f.name == "sw.js" else "max-age=300"
        self.send(200, data, TYPES.get(f.suffix, "application/octet-stream"), cache)

    def do_POST(self):
        if self.path not in ("/api/send", "/api/resolve"):
            return self.send(404, '{"error":"not found"}')
        # a custom header forces a CORS preflight, so other sites cannot post here
        if self.headers.get("X-Rig-Pocket") != "1" or (ORIGIN and self.headers.get("Origin") != ORIGIN):
            return self.send(403, '{"error":"forbidden"}')
        try:
            body = json.loads(self.rfile.read(min(int(self.headers.get("Content-Length", 0)), 8192)))
        except (ValueError, json.JSONDecodeError):
            return self.send(400, '{"error":"bad json"}')
        text = (body.get("text") or "").strip()
        if self.path == "/api/resolve":
            qid = body.get("id", "")
            if not QITEM.match(qid) or qid not in open_gates():
                return self.send(404, '{"error":"unknown decision"}')
            if not text or len(text) > 2000:
                return self.send(400, '{"error":"empty decision"}')
            out = rig("queue", "resolve", qid, "--decision", text, env=HUMAN_ENV)
            return self.send(200, json.dumps({"ok": True, "result": out.strip()[:500]}))
        seat, key = body.get("seat", ""), body.get("key")
        if not SEAT.match(seat) or seat not in known_seats():
            return self.send(404, '{"error":"unknown seat"}')
        if key in KEYS:
            keys = "\x1b" if key == "esc" else key
            out = rig("send", seat, keys, "--dangerously-interact", "--reason", "rig-pocket: answered from phone")
        elif text and len(text) <= 2000:
            out = rig("send", seat, f"From the person (via rig-pocket on their phone): {text}")
        else:
            return self.send(400, '{"error":"nothing to send"}')
        self.send(200, json.dumps({"ok": True, "result": out.strip()[:500]}))


def main():
    threading.Thread(target=poll_loop, daemon=True).start()
    ThreadingHTTPServer((BIND, PORT), Handler).serve_forever()


if __name__ == "__main__":
    main()