#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Sequential re-runner for failed indices. Usage:
   python3 rerun.py projects.json 32,51,47 shots_dir snap_dir
Like the skill's batch_screenshot_snapshot.py but:
  - accepts an explicit comma-separated index list
  - 14s settle wait (SPA / slow-iframe sites)
  - generic iframe fallback for ANY url (Qinghai / ggzy.gov.cn shells)
  - NEVER calls `close --all` (that killed sibling batches' sessions)
  - hard-fails loudly if the screenshot is a blank page (<60KB)
"""
import json, os, subprocess, sys, time

AB = ["npx", "agent-browser", "--engine", "chrome"]


def run(args, timeout=90):
    try:
        r = subprocess.run(args, capture_output=True, text=True, timeout=timeout)
        return r.stdout + r.stderr
    except Exception as e:
        return "ERR %s" % e


def pj(out):
    s = out.find("{")
    if s < 0:
        return None
    d = 0
    for i in range(s, len(out)):
        if out[i] == "{":
            d += 1
        elif out[i] == "}":
            d -= 1
            if d == 0:
                try:
                    j = json.loads(out[s:i + 1])
                    return j["data"] if isinstance(j.get("data"), dict) else j
                except Exception:
                    return None
    return None


def main():
    pfile, idxs, shot_dir, snap_dir = sys.argv[1:5]
    projs = json.load(open(pfile, encoding="utf-8"))
    ids = [int(x) for x in idxs.split(",")]
    for idx in ids:
        p = projs[idx - 1]
        url = p["url"]
        sess = "rr%02d" % idx
        png = os.path.join(shot_dir, "p%02d.png" % idx)
        txtp = os.path.join(snap_dir, "s%02d.txt" % idx)
        print("=== #%d %s ===" % (idx, p["title"][:32]), flush=True)
        run(AB + ["--session", sess, "--json", "open", url])
        time.sleep(14)
        j = pj(run(AB + ["--session", sess, "--json", "eval", "document.body.innerText"]))
        text = str(j.get("result")) if j and j.get("result") else ""
        ttl = ""
        jt = pj(run(AB + ["--session", sess, "--json", "eval", "document.title"]))
        if jt and jt.get("result"):
            ttl = str(jt["result"])
        if len(text) < 600:
            time.sleep(6)
            j3 = pj(run(AB + ["--session", sess, "--json", "eval",
                "(function(){var f=document.querySelector('iframe');if(!f)return '';"
                "try{return f.contentDocument?f.contentDocument.body.innerText:''}"
                "catch(e){return 'ERR:'+e.message}})()"]))
            if j3 and j3.get("result"):
                text += "\n[IFRAME]\n" + str(j3["result"])
        run(AB + ["--session", sess, "--json", "screenshot", "--full", png])
        if not os.path.exists(png) or os.path.getsize(png) < 60000:
            time.sleep(4)
            run(AB + ["--session", sess, "--json", "screenshot", "--full", png])
        time.sleep(2)
        run(AB + ["--session", sess, "--json", "open", url])       # liveness re-verify
        time.sleep(6)
        run(AB + ["--session", sess, "--json", "eval", "document.title"])
        run(AB + ["--session", "close", sess], timeout=30)
        head = ("项目名称：%s\n采购单位：%s\n预算金额：%s万元\n公告正文：\n%s\n发布日期：%s\n来源URL：%s\n"
                % (p["title"], p.get("buyer", ""), p.get("budget", "未披露"), text,
                   p.get("pub_date", ""), url))
        if not os.path.exists(txtp) or len(text.strip()) > 80:
            open(txtp, "w", encoding="utf-8").write(head)
        sz = os.path.getsize(png) if os.path.exists(png) else 0
        print("RESULT #%d png=%dB txt=%dch title=%s" % (idx, sz, len(text), ttl[:30]), flush=True)


main()