#!/usr/bin/env python3
"""Batch crawler scheduler — 每30分钟跑一批，自管理状态
状态文件: /tmp/crawl_batch_state.json（每天凌晨自动重置）
结果文件: /tmp/crawl_results.json（给 admin 页显示日跑状态）
"""
import json, os, sys, subprocess, time, re
from datetime import datetime, date

DB_PATH = "/root/search.db"
STATE_FILE = "/tmp/crawl_batch_state.json"
RESULTS_FILE = "/tmp/crawl_results.json"  # 给 admin 页读取日跑状态
CONFIG_FILE = "/root/gov_crawler/daily_crawl_config.json"
BATCH_SIZE = 50
TIMEOUT = 600
BATCH_TIMEOUT = 100  # 每个批次最长100秒（配合cron 120s限制）


def load_state():
    today = str(date.today())
    if os.path.exists(STATE_FILE):
        try:
            with open(STATE_FILE) as f:
                state = json.load(f)
            if state.get("date") == today:
                if state.get("completed"):
                    return state
                return state
        except (json.JSONDecodeError, KeyError):
            pass
    return {
        "date": today,
        "batch_size": BATCH_SIZE,
        "current_index": 0,
        "total": 0,
        "results": {"success": [], "failed": [], "errors": []},
        "completed": False,
        "has_new_data": False,
    }


def save_state(state):
    with open(STATE_FILE, "w") as f:
        json.dump(state, f, ensure_ascii=False, indent=2)


def save_results(batch_results):
    """Save per-script results for admin page display"""
    try:
        today = str(date.today())
        existing = []
        if os.path.exists(RESULTS_FILE):
            try:
                with open(RESULTS_FILE) as f:
                    existing = json.load(f)
                if isinstance(existing, list):
                    existing = [r for r in existing if r.get("date") == today]
                else:
                    existing = []
            except (json.JSONDecodeError, Exception):
                existing = []
        existing.extend(batch_results)
        with open(RESULTS_FILE, "w") as f:
            json.dump(existing, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"  [WARN] 保存日跑结果失败: {e}")


def load_config():
    with open(CONFIG_FILE) as f:
        cfg = json.load(f)["crawlers"]
    return [c for c in cfg if c.get("enabled", False)]


def sync_fts():
    """Sync new gov_raw records to gov_search FTS index"""
    try:
        import sqlite3
        conn = sqlite3.connect(DB_PATH)
        conn.execute('PRAGMA busy_timeout=5000')
        c = conn.cursor()
        last = c.execute("SELECT COALESCE(MAX(rowid), 0) FROM gov_search").fetchone()[0]
        c.execute("""
            INSERT OR IGNORE INTO gov_search(rowid, site_name, title, summary)
            SELECT id, site_name, title,
              substr(coalesce(nullif(content,''), summary, ''), 1, 500)
            FROM gov_raw WHERE id > ?
        """, (last,))
        synced = conn.total_changes
        if synced:
            conn.commit()
            print(f"  FTS: +{synced} rows indexed")
        conn.close()
    except Exception as e:
        print(f"  \u26a0 FTS sync: {e}")

def restart_search_app():
    try:
        subprocess.run(
            ["systemctl", "restart", "search-app.service"],
            capture_output=True, timeout=10,
        )
        print("  [服务] search-app 已重启")
    except Exception as e:
        print(f"  [服务] 重启失败: {e}")


def extract_new_count(stdout, stderr):
    text = stdout + "\n" + stderr
    for line in text.split("\n"):
        for pat in [r"新增了?(\d+)", r"增量(\d+)", r"[Nn]ew[: ]*(\d+)",
                     r"\+(\d+)\s*条", r"本次新增[了]?(\d+)"]:
            m = re.search(pat, line)
            if m:
                return int(m.group(1))
    return 0


def run_batch(state, enabled):
    start = state["current_index"]
    page = start // state["batch_size"] + 1
    end = min(start + state["batch_size"], len(enabled))
    batch = enabled[start:end]

    t_start = time.time()
    batch_has_new = False
    batch_results = []

    print(f"{'='*60}")
    print(f"  批次 {page} | {start+1}-{end}/{len(enabled)} | {len(batch)} 个爬虫")
    print(f"  时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'='*60}\n")

    for idx, c in enumerate(batch, start=start + 1):
        name = c.get("display_name", c.get("site_name", c.get("name", "Unknown")))
        script = c["script"]
        args_raw = c.get("args", "")
        if isinstance(args_raw, list):
            args = args_raw
        else:
            args = [a for a in args_raw.strip().split() if a]

        print(f"  [{idx}/{len(enabled)}] {name}...", end=" ", flush=True)

        result_entry = {
            "script": script,
            "name": name,
            "ts": time.time(),
            "date": str(date.today()),
            "has_error": True,
            "rc": -1,
            "new_count": 0,
        }

        try:
            cmd = ["python3", script]
            if args:
                cmd.extend(args)

            script_timeout = c.get("timeout", 60)
            t0 = time.time()
            r = subprocess.run(
                cmd, capture_output=True, text=True,
                timeout=script_timeout, cwd="/root/gov_crawler",
            )
            elapsed = time.time() - t0

            result_entry["ts"] = time.time()
            result_entry["rc"] = r.returncode

            if r.returncode == 0:
                new_count = extract_new_count(r.stdout, r.stderr)
                status = f"+{new_count}条" if new_count > 0 else "✓ 无新增"
                print(f"{status} ({elapsed:.0f}s)")
                state["results"]["success"].append({
                    "name": name, "new": new_count
                })
                result_entry["has_error"] = False
                result_entry["new_count"] = new_count
                if new_count > 0:
                    batch_has_new = True
                    state["has_new_data"] = True
            else:
                err = r.stderr.strip()[:100] or "exit code != 0"
                print(f"✗ FAIL ({elapsed:.0f}s) — {err}")
                state["results"]["failed"].append({
                    "name": name, "error": r.stderr[:200]
                })

        except subprocess.TimeoutExpired:
            result_entry["timeout"] = True
            print("⏰ 超时")
            state["results"]["failed"].append({
                "name": name, "error": "timeout"
            })
        except Exception as e:
            print(f"✗ ERROR: {e}")
            state["results"]["errors"].append({
                "name": name, "error": str(e)[:200]
            })

        batch_results.append(result_entry)

        # 每跑完一个就刷状态
        state["current_index"] = idx
        save_state(state)

        # 批次超时保护：超过100秒就停，剩余下次跑
        if time.time() - t_start > BATCH_TIMEOUT:
            remaining = len(batch) - (idx - start)
            print(f"  [TIMEOUT] 批次超时, {remaining} 个爬虫留待下次")
            break

    # 保存本批次结果到 admin 页可读文件
    save_results(batch_results)

    elapsed = time.time() - t_start
    s = state["results"]
    ok = len(s["success"])
    fail = len(s["failed"])
    err = len(s["errors"])
    new_total = sum(item.get("new", 0) for item in s["success"])

    print(f"\n  ─── 批次 {page} 小结 ───")
    print(f"  耗时: {elapsed:.0f}s | ✓{ok} ✗{fail} !{err}")
    print(f"  今日进度: {state['current_index']}/{len(enabled)}")
    print(f"  今日新增: {new_total} 条")

    if batch_has_new:
        restart_search_app()

    if end >= len(enabled):
        state["completed"] = True
        print(f"\n{'='*60}")
        print(f"  ✅ 今日全部完成！{len(enabled)} 个爬虫已处理")
        print(f"  新增总计: {new_total} 条")
        print(f"{'='*60}")
        if state["has_new_data"]:
            restart_search_app()

    save_state(state)


def main():
    state = load_state()

    if state.get("completed"):
        print(f"[{datetime.now().strftime('%H:%M:%S')}] 今日所有批次已完成，跳过。")
        return

    enabled = load_config()
    state["total"] = len(enabled)

    if state["current_index"] >= len(enabled):
        state["completed"] = True
        save_state(state)
        print(f"[{datetime.now().strftime('%H:%M:%S')}] 今日所有爬虫已完成！")
        return

    run_batch(state, enabled)


if __name__ == "__main__":
    main()