#!/usr/bin/env python3
"""Poll 2925 Webmail for an OpenAI/ChatGPT code matching the current generated alias.

Why this exists:
- 2925 OpenAI mail may land in 垃圾箱 instead of 收件箱.
- 垃圾箱/已删除 can contain stale OpenAI codes for older aliases.
- Never submit a visible 6-digit code unless the surrounding message matches the current alias/local-part.

Outputs compact JSONL and saves the code to --out-file without printing it.
"""
import argparse
import json
import pathlib
import re
import time
import urllib.request
import websocket
import itertools

FOLDERS = {
    "inbox": "https://2925.com/#/mailList?type=%E6%94%B6%E4%BB%B6%E7%AE%B1",
    "junk": "https://2925.com/#/mailList?type=%E5%9E%83%E5%9C%BE%E7%AE%B1",
    "deleted": "https://2925.com/#/mailList?type=%E5%B7%B2%E5%88%A0%E9%99%A4",
}


def log(obj):
    print(json.dumps({**obj, "ts": int(time.time())}, ensure_ascii=False), flush=True)


def cdp_json(base, path):
    return json.loads(urllib.request.urlopen(base + path, timeout=10).read().decode())


def connect_tab(base, url_substr="2925.com"):
    tabs = cdp_json(base, "/json")
    tab = next((t for t in tabs if url_substr in (t.get("url") or "")), tabs[0])
    ws = websocket.create_connection(tab["webSocketDebuggerUrl"], timeout=20)
    counter = itertools.count(1)

    def call(method, params=None, timeout=25):
        msg_id = next(counter)
        ws.send(json.dumps({"id": msg_id, "method": method, "params": params or {}}))
        end = time.time() + timeout
        while time.time() < end:
            msg = json.loads(ws.recv())
            if msg.get("id") == msg_id:
                if "error" in msg:
                    raise RuntimeError(msg["error"])
                return msg.get("result", {})
        raise TimeoutError(method)

    call("Runtime.enable")
    call("Page.enable")

    def eval_js(expr, await_promise=False, timeout=25):
        result = call(
            "Runtime.evaluate",
            {"expression": expr, "awaitPromise": await_promise, "returnByValue": True},
            timeout,
        )
        return result.get("result", {}).get("value")

    return call, eval_js


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cdp-port", type=int, required=True)
    ap.add_argument("--account", required=True, help="Selected 2925 base account, e.g. user@2925.com")
    ap.add_argument("--alias", required=True, help="Current generated alias used on OpenAI")
    ap.add_argument("--out-file", default="/tmp/openai_2925_email_code.txt")
    ap.add_argument("--attempts", type=int, default=15)
    ap.add_argument("--interval", type=float, default=15.0)
    ap.add_argument("--folders", default="inbox,junk", help="comma list: inbox,junk,deleted")
    ap.add_argument("--reject-code", action="append", default=[])
    args = ap.parse_args()

    base = f"http://127.0.0.1:{args.cdp_port}"
    account = args.account.strip().lower()
    alias = args.alias.strip().lower()
    alias_local = alias.split("@", 1)[0]
    rejected = {str(x).strip() for x in args.reject_code if str(x).strip()}
    folders = [x.strip() for x in args.folders.split(",") if x.strip()]

    call, ev = connect_tab(base)

    for attempt in range(1, args.attempts + 1):
        for folder in folders:
            url = FOLDERS.get(folder)
            if not url:
                continue
            call("Page.navigate", {"url": url})
            time.sleep(4)
            js = f"""
(() => {{
  try {{
    const account = {json.dumps(account)};
    const aliasLocal = {json.dumps(alias_local)};
    const txt = document.body?.innerText || '';
    const emails = [...new Set([...txt.matchAll(/[A-Za-z0-9._%+-]+@2925\\.com/g)].map(m => m[0].toLowerCase()))];
    const lines = txt.split(/\n+/).map(s => s.trim()).filter(Boolean);
    const snippets = [];
    for (let i = 0; i < lines.length; i++) {{
      const line = lines[i].toLowerCase();
      if (line.includes(aliasLocal) || /openai|chatgpt|tm1\\.openai|验证码|verification/.test(line)) {{
        const s = lines.slice(Math.max(0, i - 5), Math.min(lines.length, i + 9)).join('\n');
        if (s.toLowerCase().includes(aliasLocal) && /OpenAI|ChatGPT|tm1\\.openai|验证码|verification/i.test(s)) snippets.push(s);
      }}
    }}
    const codes = [...new Set(snippets.flatMap(s => [...s.matchAll(/(?<!\\d)(\\d{{6}})(?!\\d)/g)].map(m => m[1])))];
    return {{ href: location.href, emails, hasAccount: emails.includes(account), hasAliasLocal: txt.toLowerCase().includes(aliasLocal), snippets: snippets.slice(0, 5), codes }};
  }} catch (e) {{
    return {{ error: String(e), emails: [], snippets: [], codes: [] }};
  }}
}})()
"""
            data = ev(js) or {"error": "null_eval", "emails": [], "snippets": [], "codes": []}
            if account not in data.get("emails", []):
                log({"phase": "account_gate_failed", "folder": folder, "expected": account, "emails": data.get("emails", [])})
                raise SystemExit(4)
            fresh = [c for c in data.get("codes", []) if c not in rejected]
            log({
                "phase": "poll_alias_code",
                "attempt": attempt,
                "folder": folder,
                "hasAliasLocal": bool(data.get("hasAliasLocal")),
                "snippetCount": len(data.get("snippets", [])),
                "freshCount": len(fresh),
            })
            if fresh:
                pathlib.Path(args.out_file).write_text(fresh[0])
                log({"phase": "code_saved", "folder": folder, "attempt": attempt, "outFile": args.out_file})
                return
        time.sleep(args.interval)

    log({"phase": "timeout_no_current_alias_code", "alias": alias})
    raise SystemExit(2)


if __name__ == "__main__":
    main()