#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""增量重建 gov_bigram v2: LEFT JOIN 快速找缺失, 分批补 2-gram"""
import sqlite3, re, time

DB = '/root/search.db'
conn = sqlite3.connect(DB, timeout=120)
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA busy_timeout=120000")
c = conn.cursor()

# 快速找缺失: gov_raw LEFT JOIN bigram(distinct rowid) 取 NULL
print("查找缺失记录...", flush=True)
t0 = time.time()
missing = c.execute('''
    SELECT r.id, r.title, COALESCE(r.content, ''), COALESCE(r.summary, '')
    FROM gov_raw r
    LEFT JOIN (SELECT DISTINCT rowid FROM gov_bigram) b ON b.rowid = r.id
    WHERE b.rowid IS NULL
''').fetchall()
print(f"缺失: {len(missing)} 条, 查找耗时 {time.time()-t0:.0f}s", flush=True)

def extract_grams(text):
    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):
    """写 bigram 批次, 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(8 * (attempt + 1))
                continue
            raise
    return False

t0 = time.time()
done = 0
batch = []
for rid, title, content, summary in missing:
    grams = extract_grams(title) | extract_grams(content) | extract_grams(summary)
    if not grams:
        grams = {'__empty__'}
    for g in grams:
        batch.append((rid, g))
    done += 1
    if len(batch) >= 50000:
        insert_batch(batch)
        batch = []
        print(f"  {done}/{len(missing)} 条, {time.time()-t0:.0f}s", flush=True)

if batch:
    insert_batch(batch)

elapsed = time.time() - t0
total = c.execute('SELECT COUNT(*) FROM gov_bigram').fetchone()[0]
print(f"完成: {done} 条, 耗时 {elapsed:.0f}s, gov_bigram 总行数 {total}", flush=True)
conn.close()