#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""增量重建 gov_bigram —— 水位线版（2026-09-11 替换 inc2）

为什么换掉 inc2：
  旧 inc2 用 `LEFT JOIN (SELECT DISTINCT rowid FROM gov_bigram) b ... WHERE b.rowid IS NULL`
    ① `SELECT DISTINCT rowid FROM gov_bigram` 全扫 1.42 亿条索引项 —— 实测**仅"查找缺失"就 82s**
       （即使 0 条缺失也照跑）
    ② 结果再 `.fetchall()` 全部读进内存 —— 生产机只有 1.75GB 内存，数据量一大就有 OOM 风险

新写法（两处关键优化）：
  ① **meta 水位线表** `index_watermark(name, value)`，O(1) 读。
     ⚠️ 不能直接 `SELECT MAX(rowid) FROM gov_bigram`：本表 `rowid` 是**显式 INTEGER 列**
     （见 pragma_table_info），没有以它打头的索引 → MAX() 退化为全扫 1.42 亿行，
     实测 **38s**（每次运行都要付，哪怕 0 条新数据）。
     `MAX(_rowid_)` 虽 0.003s 但那是插入序号，**不是 doc id**，不可用。
  ② 只取 `gov_raw.id > 水位线` 的行，keyset 分批 200 条流式处理，
     峰值内存 = 一个批次的 2-gram（几百 KB）。
  ③ 与 inc2 完全同口径的 gram 提取（连续中文/数字/字母段内所有 2-gram）
  ④ 崩溃后重跑安全：按 id 升序写，水位线只前进，不重复不漏。

⚠️ 前提：增量写入按 id 升序（本脚本自身保证 keyset ORDER BY id）。
   历史遗留的"水位线以下缺失"用 --verify 一次性全扫核对（只读）。

用法:
  python3 rebuild_bigram_inc2.py            # 日常增量（cron + 日跑收尾均调用）
  python3 rebuild_bigram_inc2.py --verify   # 额外做一次权威全扫核对（慢，平时勿用）
  python3 rebuild_bigram_inc2.py --limit=500
"""
import sqlite3, re, sys, time

DB = "/root/search.db"   # 软链 → /mnt/data/search.db
WM_TABLE = "index_watermark"
WM_NAME = "bigram"
VERIFY = "--verify" in sys.argv
LIMIT = None
for _a in sys.argv:
    if _a.startswith("--limit="):
        LIMIT = int(_a.split("=", 1)[1])

BATCH = 200          # 每批行数（内存上界）
GRAM_FLUSH = 50000   # 每批待写 gram 数上限


def log(msg):
    print("[%s] %s" % (time.strftime("%Y-%m-%d %H:%M:%S"), msg), flush=True)


def extract_grams(text):
    """与 inc2 完全同口径：连续 [中文/数字/字母] 段内所有 2-gram"""
    if not text:
        return set()
    t = text.lower()
    grams = set()
    for seg in re.findall(r"[\u4e00-\u9fff0-9a-z]+", t):
        if len(seg) < 2:
            continue
        for i in range(len(seg) - 1):
            grams.add(seg[i:i + 2])
    return grams


def insert_batch(c, conn, batch, retries=40):
    """写批次；database is locked 时退避重试（调度器并行写库时防崩）"""
    for attempt in range(retries):
        try:
            c.executemany("INSERT OR IGNORE INTO gov_bigram (rowid, gram) VALUES (?,?)", batch)
            conn.commit()
            return True
        except sqlite3.OperationalError as e:
            if "locked" in str(e).lower() and attempt < retries - 1:
                time.sleep(5 * (attempt + 1))
                continue
            raise
    return False


def get_watermark(c, conn):
    """读水位线；首次运行用 MAX(rowid) 初始化（一次性 38s）并存表"""
    c.execute("CREATE TABLE IF NOT EXISTS %s(name TEXT PRIMARY KEY, value INTEGER)" % WM_TABLE)
    conn.commit()
    row = c.execute("SELECT value FROM %s WHERE name=?" % WM_TABLE, (WM_NAME,)).fetchone()
    if row:
        return row[0]
    log("首次运行：初始化水位线（MAX(rowid) 全扫，一次性较慢）...")
    t0 = time.time()
    wm = c.execute("SELECT COALESCE(MAX(rowid), 0) FROM gov_bigram").fetchone()[0]
    c.execute("INSERT OR REPLACE INTO %s(name, value) VALUES(?,?)" % WM_TABLE, (WM_NAME, wm))
    conn.commit()
    log("初始化完成 wm=%d (%.0fs)" % (wm, time.time() - t0))
    return wm


def set_watermark(c, conn, value):
    c.execute("INSERT OR REPLACE INTO %s(name, value) VALUES(?,?)" % WM_TABLE, (WM_NAME, value))
    conn.commit()


def main():
    conn = sqlite3.connect(DB, timeout=120)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=120000")
    conn.execute("PRAGMA synchronous=NORMAL")
    c = conn.cursor()

    wm = get_watermark(c, conn)
    gap = c.execute("SELECT COUNT(*) FROM gov_raw WHERE id > ?", (wm,)).fetchone()[0]
    log("水位线=%d | 待补 %d 行" % (wm, gap))

    if gap == 0 and not VERIFY:
        log("无新行，跳过（0 写入）")
        conn.close()
        return

    last, done, grams_done = wm, 0, 0
    t0 = time.time()
    if gap:
        while True:
            rows = c.execute(
                "SELECT id, title, COALESCE(content,''), COALESCE(summary,'') "
                "FROM gov_raw WHERE id > ? ORDER BY id LIMIT ?", (last, BATCH)).fetchall()
            if not rows:
                break
            pending = []
            for rid, title, content, summary in rows:
                g = extract_grams(title) | extract_grams(content) | extract_grams(summary)
                if not g:
                    g = {"__empty__"}
                pending.extend((rid, x) for x in g)
                if len(pending) >= GRAM_FLUSH:
                    insert_batch(c, conn, pending)
                    grams_done += len(pending)
                    pending = []
            if pending:
                insert_batch(c, conn, pending)
                grams_done += len(pending)
            last = rows[-1][0]
            done += len(rows)
            set_watermark(c, conn, last)          # 每批推进，崩溃可续
            if done % 1000 < BATCH:
                log("  %d/%d 条, %d gram, %.0fs" % (done, gap, grams_done, time.time() - t0))
            if LIMIT and done >= LIMIT:
                log("(--limit 命中，停止)")
                break

    remain = c.execute("SELECT COUNT(*) FROM gov_raw WHERE id > ?", (last,)).fetchone()[0]
    log("完成: %d 条, %d gram, %.0fs | 水位线 -> %d | 剩余缺口 %d"
        % (done, grams_done, time.time() - t0, last, remain))

    if VERIFY:
        log("--verify: 全扫核对（较慢）...")
        tv = time.time()
        miss = c.execute("""
            SELECT COUNT(*) FROM gov_raw r
            LEFT JOIN (SELECT DISTINCT rowid FROM gov_bigram) b ON b.rowid = r.id
            WHERE b.rowid IS NULL
        """).fetchone()[0]
        log("--verify 结果: 水位线以下遗漏 %d 条 (%.0fs)" % (miss, time.time() - tv))
    conn.close()


if __name__ == "__main__":
    main()