#!/usr/bin/env python3
"""Patch: guest/user/owner plans with server-side session limits,
owner admin access, keyless host registration support, host/user
management endpoints, public host download."""
import re

p = "server/netdesk_server.py"
s = open(p, encoding="utf-8").read()

# ---------- 1. plans + owner ----------
old = "TUNNEL_TID = [1]"
new = """TUNNEL_TID = [1]

# session plans: guests 20 min, verified users 45 min, owner unlimited
PLAN_LIMITS = {"owner": None, "user": 45 * 60, "guest": 20 * 60}

def plan_for(email):
    if email and ARGS.owner_email and email.lower() == ARGS.owner_email.lower():
        return "owner", None
    if email:
        return "user", PLAN_LIMITS["user"]
    return "guest", PLAN_LIMITS["guest"]"""
assert old in s, "p1"
s = s.replace(old, new, 1)

# ---------- 2. owner + plan on session.begin; drop hard signin gate ----------
old = """                elif t == "session.begin":
                    if ARGS.mail_from and not user_from_cookie(req):
                        log_event("auth", "connect blocked: viewer not signed in")
                        await send_json(ws, {"t": "session.deny", "sid": None, "reason": "signin"})
                        continue
                    hid = str(m.get("id", "")).strip()
                    h = HOSTS.get(hid)"""
new = """                elif t == "session.begin":
                    vuser = user_from_cookie(req)
                    vplan, vlimit = plan_for(vuser)
                    hid = str(m.get("id", "")).strip()
                    h = HOSTS.get(hid)"""
assert old in s, "p2"
s = s.replace(old, new, 1)

# after sess created, attach plan info
old = """                    sess = Session(sid, h, ws)
                    h.session = sess
                    viewer_sid = sid
                    STATS.sessions_total += 1
                    log_event("session", f"viewer {peer} connecting to host {hid}")
                    log.info("session %s: viewer %s -> host %s", sid, peer, hid)
                    await send_json(h.ws, {"t": "session.begin", "sid": sid})"""
new = """                    sess = Session(sid, h, ws)
                    sess.plan, sess.vlimit, sess.vemail = vplan, vlimit, vuser
                    h.session = sess
                    viewer_sid = sid
                    STATS.sessions_total += 1
                    log_event("session", f"viewer {peer} ({vplan}) connecting to host {hid}")
                    log.info("session %s: viewer %s (%s) -> host %s", sid, peer, vplan, hid)
                    await send_json(h.ws, {"t": "session.begin", "sid": sid})"""
assert old in s, "p3"
s = s.replace(old, new, 1)

# on accept: tell viewer the plan + start the enforcement timer
old = """                    if t == "session.accept" and host is not None:
                        sess.active = True
                        log_event("session", f"session accepted on host {sess.host.id}")
                        await send_json(sess.viewer_ws, m)"""
new = """                    if t == "session.accept" and host is not None:
                        sess.active = True
                        log_event("session", f"session accepted on host {sess.host.id}")
                        plan_msg = {"t": "session.plan", "plan": sess.plan,
                                    "limit": sess.vlimit or 0, "email": sess.vemail or ""}
                        await send_json(sess.viewer_ws, plan_msg)
                        if sess.vlimit:
                            async def _limit_killer(s=sess):
                                await asyncio.sleep(s.vlimit)
                                if s.active:
                                    log_event("session", f"time limit on host {s.host.id}")
                                    await relay_session_end(
                                        s, "session time limit reached - reconnect or sign up",
                                        None)
                                    if s.host:
                                        s.host.session = None
                            asyncio.ensure_future(_limit_killer(sess))
                        await send_json(sess.viewer_ws, m)"""
assert old in s, "p4"
s = s.replace(old, new, 1)

# ---------- 3. Session fields ----------
old = """class Session:
    def __init__(self, sid, host, viewer_ws):
        self.sid = sid
        self.host = host
        self.viewer_ws = viewer_ws
        self.active = False"""
new = """class Session:
    def __init__(self, sid, host, viewer_ws):
        self.sid = sid
        self.host = host
        self.viewer_ws = viewer_ws
        self.active = False
        self.plan = "guest"
        self.vlimit = None
        self.vemail = ""
        self.started = time.time()"""
assert old in s, "p5"
s = s.replace(old, new, 1)

# ---------- 4. owner-aware portal auth + admin endpoints ----------
old = """def portal_authorized(req) -> bool:
    if not ARGS.portal_password:
        return False
    tok = req.cookies.get("ndportal", "")
    return tok in PORTAL_SESSIONS"""
new = """def owner_from_cookie(req) -> bool:
    email = user_from_cookie(req)
    return bool(email and ARGS.owner_email and email.lower() == ARGS.owner_email.lower())


def portal_authorized(req) -> bool:
    if owner_from_cookie(req):
        return True
    if not ARGS.portal_password:
        return False
    tok = req.cookies.get("ndportal", "")
    return tok in PORTAL_SESSIONS"""
assert old in s, "p6"
s = s.replace(old, new, 1)

# admin endpoints
old = """async def portal_download(req):"""
new = """async def portal_users(req):
    if not portal_authorized(req):
        return web.json_response({"error": "unauthorized"}, status=401)
    users = []
    with _acct_db() as conn:
        for row in conn.execute("SELECT email, verified, created FROM users "
                                "ORDER BY created DESC LIMIT 500"):
            users.append({"email": row[0], "verified": bool(row[1]), "created": row[2]})
    return web.json_response({"users": users})


async def portal_user_delete(req):
    if not portal_authorized(req):
        return web.json_response({"error": "unauthorized"}, status=401)
    body = await req.json()
    email = str(body.get("email", "")).lower()
    if ARGS.owner_email and email == ARGS.owner_email.lower():
        return web.json_response({"error": "cannot delete the owner"}, status=400)
    with _acct_db() as conn:
        conn.execute("DELETE FROM users WHERE email=?", (email,))
    for tok, em in list(USER_SESSIONS.items()):
        if em == email:
            USER_SESSIONS.pop(tok, None)
    log_event("admin", f"user deleted: {email}")
    return web.json_response({"ok": True})


async def portal_host_delete(req):
    if not portal_authorized(req):
        return web.json_response({"error": "unauthorized"}, status=401)
    hid = str((await req.json()).get("id", ""))
    h = HOSTS.get(hid)
    if h is None:
        return web.json_response({"error": "host offline (removed from registry)"}, status=404)
    try:
        await h.ws.close()
    except Exception:
        pass
    log_event("admin", f"host {hid} removed by admin")
    return web.json_response({"ok": True, "msg": "host removed; it must re-register"})


async def portal_session_disconnect(req):
    if not portal_authorized(req):
        return web.json_response({"error": "unauthorized"}, status=401)
    hid = str((await req.json()).get("id", ""))
    h = HOSTS.get(hid)
    if h is None or not h.session:
        return web.json_response({"error": "no active session"}, status=404)
    await relay_session_end(h.session, "disconnected by admin", None)
    h.session = None
    log_event("admin", f"session on {hid} disconnected by admin")
    return web.json_response({"ok": True})


async def portal_download(req):"""
assert old in s, "p7"
s = s.replace(old, new, 1)

# overview: users + richer host info
old = """    return web.json_response({
        "version": VERSION,
        "reg_key": ARGS.reg_key or "",
        "uptime": int(now - STATS.started),"""
new = """    users = []
    with _acct_db() as conn:
        for row in conn.execute("SELECT email, verified, created FROM users "
                                "ORDER BY created DESC LIMIT 500"):
            users.append({"email": row[0], "verified": bool(row[1]), "created": row[2]})
    for k in known:
        h = HOSTS.get(k["id"])
        k["plan"] = h.session.plan if (h and h.session) else ""
    return web.json_response({
        "version": VERSION,
        "owner": bool(ARGS.owner_email and user_from_cookie(req) and
                      user_from_cookie(req).lower() == ARGS.owner_email.lower()),
        "users": users,
        "reg_key": ARGS.reg_key or "",
        "uptime": int(now - STATS.started),"""
assert old in s, "p8"
s = s.replace(old, new, 1)

# routes
old = """    app.router.add_get("/portal/api/overview", portal_overview)"""
new = """    app.router.add_get("/portal/api/overview", portal_overview)
    app.router.add_get("/portal/api/users", portal_users)
    app.router.add_post("/portal/api/user/delete", portal_user_delete)
    app.router.add_post("/portal/api/host/delete", portal_host_delete)
    app.router.add_post("/portal/api/session/disconnect", portal_session_disconnect)"""
assert old in s, "p9"
s = s.replace(old, new, 1)

# public host download (homepage button)
old = """async def tunnel_ws(req):"""
new = """async def public_download(req):
    name = os.path.basename(req.match_info.get("name", ""))
    path = os.path.join(FILES_DIR, name)
    if not os.path.isfile(path):
        return web.Response(status=404, text="not found")
    with open(path, "rb") as f:
        body = f.read()
    resp = web.Response(body=body, content_type="application/octet-stream")
    resp.headers["Content-Disposition"] = f'attachment; filename="{name}"'
    return resp


async def tunnel_ws(req):"""
assert old in s, "p10"
s = s.replace(old, new, 1)
old = """    app.router.add_get("/tunnel/{hid}", tunnel_ws)"""
new = """    app.router.add_get("/tunnel/{hid}", tunnel_ws)
    app.router.add_get("/download/{name}", public_download)"""
assert old in s; s = s.replace(old, new, 1)

# ---------- 5. CLI owner arg ----------
old = """    ap.add_argument("--mail-from", dest="mail_from", default="","""
new = """    ap.add_argument("--owner-email", dest="owner_email", default="",
                    help="email with unrestricted access + admin panel")
    ap.add_argument("--mail-from", dest="mail_from", default="","""
assert old in s; s = s.replace(old, new, 1)

open(p, "w", encoding="utf-8").write(s)
print("plans/owner/admin patch applied")