#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""bigram 水位线增量补录（内存安全版）

与 rebuild_bigram_inc2.py 的差异：
  旧: SELECT DISTINCT rowid FROM gov_bigram（8600 万行全扫）+
      LEFT JOIN + fetchall()（把全部缺失行读进内存）
  新: 只取 max(rowid) 水位线之上的行，keyset 分批 200 条流式处理，
      峰值内存 = 一个批次的 gram（几百 KB）
用法: python3 bigram_backfill_wm.py [--limit N]   (--limit 仅调试用)
"""
import sqlite3, re, sys, time

DB = "/root/search.db"
LIMIT = None
for _a in sys.argv:
    if _a.startswith("--limit="):
        LIMIT = int(_a.split("=", 1)[1])

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 = c.execute("SELECT COALESCE(MAX(rowid), 0) FROM gov_bigram").fetchone()[0]
gap = c.execute("SELECT COUNT(*) FROM gov_raw WHERE id > ?", (wm,)).fetchone()[0]
print("水位线 max(rowid)=%d | 待补 %d 行" % (wm, gap), flush=True)


def extract_grams(text):
    """与 rebuild_bigram_inc2.py 完全同口径：连续中文/数字/字母段内所有 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(batch, retries=40):
    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


last, done, total_grams = wm, 0, 0
t0 = time.time()
while True:
    rows = c.execute(
        "SELECT id, title, COALESCE(content,''), COALESCE(summary,'') "
        "FROM gov_raw WHERE id > ? ORDER BY id LIMIT 200", (last,)).fetchall()
    if not rows:
        break
    batch = []
    for rid, title, content, summary in rows:
        grams = extract_grams(title) | extract_grams(content) | extract_grams(summary)
        if not grams:
            grams = {"__empty__"}
        batch.extend((rid, g) for g in grams)
    insert_batch(batch)
    last = rows[-1][0]
    done += len(rows)
    total_grams += len(batch)
    if done % 1000 < 200:
        print("  %d/%d 条, %d gram, %.0fs" % (done, gap, total_grams, time.time() - t0), flush=True)
    if LIMIT and done >= LIMIT:
        print("(--limit 命中，停止)", flush=True)
        break

elapsed = time.time() - t0
new_wm = c.execute("SELECT COALESCE(MAX(rowid),0) FROM gov_bigram").fetchone()[0]
remain = c.execute("SELECT COUNT(*) FROM gov_raw WHERE id > ?", (new_wm,)).fetchone()[0]
print("完成: %d 条, %d gram, %.0fs" % (done, total_grams, elapsed), flush=True)
print("水位线 -> %d | 剩余缺口 %d" % (new_wm, remain), flush=True)
conn.close()