#!/usr/bin/env python3
"""
GHL OAuth callback handler for Rob's marketplace app ("Hermes App V2").

Catches GET /ghl/oauth/callback?code=...&state=..., exchanges the code for
access+refresh tokens at GHL, stores them at /root/.hermes/secrets/ghl-oauth.json
(chmod 600), and refreshes access tokens in the background before expiry.

MULTI-TOKEN STORE (added 2026-08-03):
  ghl-oauth.json layout:
    {
      "tokens": {
        "<locationId>":          {access_token, refresh_token, expires_at, ...},
        "agency:<companyId>":    {access_token, refresh_token, expires_at, ...}
      },
      "active": "<last-saved key>",
      ...top-level mirror of the active token payload (backward compat)...
    }
  - Tokens are keyed so a new install NEVER clobbers an existing location's token.
  - The top level mirrors the most recently saved token's fields (plus "tokens"
    and "active") so legacy readers doing json.load(...)['access_token'] keep
    working.
  - Legacy single-token files are migrated on first load (folded into "tokens").

Endpoints:
  GET /ghl/health                    — status + per-token locationId/expiry list
  GET /ghl/token?locationId=<id>     — raw token payload for that location
                                       (pass locationId=agency:<companyId> for
                                       an agency token). Omit to get the active one.
  GET /ghl/oauth/callback            — GHL redirect target

Config: /root/.hermes/secrets/ghl-oauth.env
  GHL_CLIENT_ID=...
  GHL_CLIENT_SECRET=...
  GHL_REDIRECT_URI=https://robblake.cloud/ghl/oauth/callback   (must match app registration EXACTLY)

Runs on 127.0.0.1:9120 behind nginx. systemd unit: ghl-oauth.service
"""

import json
import os
import stat
import subprocess
import threading
import time
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

ENV_FILE = "/root/.hermes/secrets/ghl-oauth.env"
TOKEN_FILE = "/root/.hermes/secrets/ghl-oauth.json"
GHL_TOKEN_URL = "https://services.leadconnectorhq.com/oauth/token"
API_VERSION = "2021-07-28"
LISTEN = ("127.0.0.1", 9120)

# Refresh when fewer than this many seconds remain before expiry
REFRESH_MARGIN = 3600
REFRESH_CHECK_INTERVAL = 600

STORE_LOCK = threading.Lock()


def load_env(path):
    cfg = {}
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            cfg[k.strip()] = v.strip()
    return cfg


CFG = load_env(ENV_FILE)
CLIENT_ID = CFG["GHL_CLIENT_ID"]
CLIENT_SECRET = CFG["GHL_CLIENT_SECRET"]
REDIRECT_URI = CFG["GHL_REDIRECT_URI"]


def log(msg):
    print(f"[ghl-oauth] {msg}", flush=True)


def post_form(url, data):
    body = urllib.parse.urlencode(data)
    cmd = [
        "curl", "-s", "-X", "POST", url,
        "-H", "Content-Type: application/x-www-form-urlencoded",
        "-H", "Accept: application/json",
        "-d", body,
        "--max-time", "30",
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=35)
    if result.returncode != 0:
        raise RuntimeError(f"curl failed: {result.stderr}")
    return json.loads(result.stdout)


def token_key(payload):
    """Storage key for a token payload: locationId, or agency:<companyId>."""
    loc = payload.get("locationId")
    if loc:
        return loc
    comp = payload.get("companyId")
    if comp:
        return f"agency:{comp}"
    return "unknown"


def _migrated(store):
    """Ensure store has the multi-token layout. Folds legacy top-level tokens in."""
    if not isinstance(store, dict):
        return {"tokens": {}, "active": None}
    if "tokens" in store and isinstance(store["tokens"], dict):
        store.setdefault("active", None)
        return store
    # Legacy layout: the store IS a single token payload
    if store.get("access_token"):
        key = token_key(store)
        payload = {k: v for k, v in store.items()}
        log(f"migrating legacy single-token file into multi-token store (key={key})")
        return {"tokens": {key: payload}, "active": key}
    return {"tokens": {}, "active": None}


def load_store():
    """Load the multi-token store, migrating legacy format if needed."""
    if not os.path.exists(TOKEN_FILE):
        return {"tokens": {}, "active": None}
    try:
        with open(TOKEN_FILE) as f:
            raw = json.load(f)
    except Exception as e:
        log(f"WARNING: could not parse token file: {e}")
        return {"tokens": {}, "active": None}
    return _migrated(raw)


def write_store(store):
    """Persist store atomically (600 perms) with the active token mirrored top-level."""
    out = {}
    active_key = store.get("active")
    active_tok = store["tokens"].get(active_key) if active_key else None
    if active_tok:
        out.update(active_tok)  # backward-compat top-level mirror
    out["tokens"] = store["tokens"]
    out["active"] = active_key
    tmp = TOKEN_FILE + ".tmp"
    with open(tmp, "w") as f:
        json.dump(out, f, indent=2)
    os.chmod(tmp, stat.S_IRUSR | stat.S_IWUSR)
    os.replace(tmp, TOKEN_FILE)


def save_tokens(payload):
    """Add/replace the token for its location key; mark it active. Never clobbers
    other locations' tokens."""
    payload = dict(payload)
    payload["obtained_at"] = int(time.time())
    if "expires_in" in payload:
        payload["expires_at"] = payload["obtained_at"] + int(payload["expires_in"])
    key = token_key(payload)
    with STORE_LOCK:
        store = load_store()
        store["tokens"][key] = payload
        store["active"] = key
        write_store(store)
    log(f"tokens saved: key={key} locationId={payload.get('locationId')} "
        f"companyId={payload.get('companyId')} userType={payload.get('userType')} "
        f"scope={str(payload.get('scope'))[:80]}")


def exchange_code(code):
    return post_form(GHL_TOKEN_URL, {
        "client_id": CLIENT_ID,
        "client_secret": CLIENT_SECRET,
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": REDIRECT_URI,
    })


def refresh_tokens(refresh_token):
    return post_form(GHL_TOKEN_URL, {
        "client_id": CLIENT_ID,
        "client_secret": CLIENT_SECRET,
        "grant_type": "refresh_token",
        "refresh_token": refresh_token,
    })


def refresh_loop():
    """Background thread: keep every stored access token fresh."""
    while True:
        time.sleep(REFRESH_CHECK_INTERVAL)
        try:
            with STORE_LOCK:
                store = load_store()
            for key, tok in list(store["tokens"].items()):
                expires_at = tok.get("expires_at", 0)
                remaining = expires_at - time.time()
                if remaining > REFRESH_MARGIN:
                    continue
                rt = tok.get("refresh_token")
                if not rt:
                    log(f"WARNING: token {key} near expiry but no refresh_token stored")
                    continue
                log(f"token {key} expires in {int(remaining)}s — refreshing")
                new = refresh_tokens(rt)
                # Preserve the key even if GHL omits locationId on refresh payloads
                new.setdefault("locationId", tok.get("locationId"))
                new.setdefault("companyId", tok.get("companyId"))
                new.setdefault("userType", tok.get("userType"))
                new = dict(new)
                new["obtained_at"] = int(time.time())
                if "expires_in" in new:
                    new["expires_at"] = new["obtained_at"] + int(new["expires_in"])
                with STORE_LOCK:
                    store = load_store()
                    store["tokens"][key] = new
                    write_store(store)  # does NOT change "active"
                log(f"refresh OK: {key}")
        except Exception as e:
            log(f"refresh error: {e}")


WEBHOOK_LOG = "/root/.hermes/ghl-oauth/webhooks.jsonl"
WEBHOOK_LOG_MAX_BYTES = 5 * 1024 * 1024  # rotate at 5MB


def record_webhook(payload, headers_note=""):
    """Append a webhook event to the JSONL log, rotating when oversized."""
    entry = {
        "received_at": int(time.time()),
        "type": payload.get("type") if isinstance(payload, dict) else None,
        "payload": payload,
    }
    if headers_note:
        entry["note"] = headers_note
    try:
        if os.path.exists(WEBHOOK_LOG) and os.path.getsize(WEBHOOK_LOG) > WEBHOOK_LOG_MAX_BYTES:
            os.replace(WEBHOOK_LOG, WEBHOOK_LOG + ".1")
    except Exception:
        pass
    with open(WEBHOOK_LOG, "a") as f:
        f.write(json.dumps(entry) + "\n")
    os.chmod(WEBHOOK_LOG, stat.S_IRUSR | stat.S_IWUSR)


class Handler(BaseHTTPRequestHandler):
    server_version = "ghl-oauth/2.1"

    def _send(self, code, body, ctype="application/json"):
        data = body.encode()
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_POST(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path != "/ghl/webhook":
            self._send(404, json.dumps({"error": "not found"}))
            return
        try:
            length = int(self.headers.get("Content-Length", 0))
        except ValueError:
            length = 0
        raw = self.rfile.read(min(length, 1024 * 1024)) if length else b""
        try:
            payload = json.loads(raw.decode("utf-8", errors="replace"))
        except Exception:
            payload = {"_raw": raw.decode("utf-8", errors="replace")[:2000]}
        etype = payload.get("type") if isinstance(payload, dict) else None
        loc = payload.get("locationId") if isinstance(payload, dict) else None
        log(f"webhook received: type={etype} locationId={loc}")
        try:
            record_webhook(payload)
        except Exception as e:
            log(f"webhook log write failed: {e}")
        # GHL expects a 200 quickly; no processing beyond logging for now
        self._send(200, json.dumps({"status": "received"}))

    def do_GET(self):
        parsed = urllib.parse.urlparse(self.path)
        qs = urllib.parse.parse_qs(parsed.query)

        if parsed.path == "/ghl/health":
            store = load_store()
            now = time.time()
            state = {
                "status": "ok",
                "token_file": os.path.exists(TOKEN_FILE),
                "active": store.get("active"),
                "tokens": {
                    key: {
                        "locationId": tok.get("locationId"),
                        "userType": tok.get("userType"),
                        "expires_in_s": int(tok.get("expires_at", 0) - now),
                    }
                    for key, tok in store["tokens"].items()
                },
            }
            self._send(200, json.dumps(state))
            return

        if parsed.path == "/ghl/token":
            store = load_store()
            key = qs.get("locationId", [store.get("active")])[0]
            tok = store["tokens"].get(key)
            if not tok:
                self._send(404, json.dumps({
                    "error": "no token for key",
                    "key": key,
                    "available": list(store["tokens"].keys()),
                }))
                return
            self._send(200, json.dumps(tok))
            return

        if parsed.path.startswith("/ghl/images/linkedin/"):
            # Serve LinkedIn post images
            filename = parsed.path.split("/")[-1]
            filepath = f"/root/.hermes/images/linkedin/{filename}"
            if os.path.exists(filepath) and filename.endswith(".png"):
                with open(filepath, "rb") as f:
                    data = f.read()
                self.send_response(200)
                self.send_header("Content-Type", "image/png")
                self.send_header("Content-Length", str(len(data)))
                self.send_header("Cache-Control", "public, max-age=31536000")
                self.end_headers()
                self.wfile.write(data)
                return
            self._send(404, json.dumps({"error": "not found"}))
            return

        if parsed.path != "/ghl/oauth/callback":
            self._send(404, json.dumps({"error": "not found"}))
            return

        err = qs.get("error", [None])[0]
        if err:
            log(f"oauth error from GHL: {err} {qs.get('error_description')}")
            self._send(400, json.dumps({"error": err,
                                        "detail": qs.get("error_description", [""])[0]}))
            return

        code = qs.get("code", [None])[0]
        if not code:
            self._send(400, json.dumps({"error": "missing code parameter"}))
            return

        log("received authorization code — exchanging")
        try:
            tokens = exchange_code(code)
            save_tokens(tokens)
        except Exception as e:
            detail = ""
            resp = getattr(e, "read", None)
            if callable(resp):
                try:
                    detail = resp().decode()[:500]
                except Exception:
                    detail = ""
            log(f"exchange failed: {e} {detail}")
            self._send(502, json.dumps({"error": "token exchange failed",
                                        "detail": f"{e} {detail}".strip()}))
            return

        key = token_key(tokens)
        self._send(200, f"""<!doctype html><html><head><title>GHL connected</title></head>
<body style="font-family:system-ui;max-width:560px;margin:80px auto;text-align:center">
<h1>&#9989; App installed</h1>
<p>Token stored for <strong>{key}</strong>. It will refresh automatically and will NOT
overwrite tokens for other locations.</p>
<p>You can close this tab.</p>
</body></html>""", "text/html")

    def log_message(self, fmt, *args):
        log("http " + (fmt % args))


if __name__ == "__main__":
    # Migrate legacy token file up-front so nothing is lost between restarts
    if os.path.exists(TOKEN_FILE):
        with STORE_LOCK:
            store = load_store()
            write_store(store)
        log(f"store ready: {list(store['tokens'].keys())} active={store.get('active')}")
    threading.Thread(target=refresh_loop, daemon=True).start()
    log(f"listening on {LISTEN[0]}:{LISTEN[1]}")
    ThreadingHTTPServer(LISTEN, Handler).serve_forever()