#!/usr/bin/env python3
"""Server-side JSONL importer v2: 保留表格结构(转|分隔行)与链接URL"""
import json, sqlite3, sys, os, re, html

def strip_html(text):
    """正文清洗: 保留 <table> 结构转 | 分隔行, a 链接转 [文本](URL), 其余转段落"""
    if not text:
        return ""
    text = html.unescape(text)
    # 1. 保护表格块
    tables = []
    def _hold_table(m):
        tables.append(m.group(0))
        return f'@@TABLE{len(tables)-1}@@'
    text = re.sub(r'<table[^>]*>.*?</table>', _hold_table, text, flags=re.DOTALL | re.IGNORECASE)
    # 2. a 链接转 [文本](URL)
    def _hold_link(m):
        txt = re.sub(r'<[^>]+>', '', m.group(2)).strip()
        href = m.group(1)
        return f'[{txt}]({href})' if txt else href
    text = re.sub(r'<a[^>]+href=["\']([^"\']+)["\'][^>]*>(.*?)</a>', _hold_link, text, flags=re.DOTALL | re.IGNORECASE)
    # 3. 其余标签: 块级转换行, 行内剥掉
    text = re.sub(r'</?(?:p|div|h[1-6]|li|br\s*/?|blockquote|section|article|tr)[^>]*>', chr(10), text, flags=re.IGNORECASE)
    text = re.sub(r'<[^>]+>', '', text)
    text = re.sub(r'[ \t]+', ' ', text)
    text = re.sub(r'\n{3,}', chr(10)*2, text)
    # 4. 表格恢复: 每行 td/th 之间用 | 分隔
    def _restore_table(t):
        rows = re.findall(r'<tr[^>]*>([\s\S]*?)</tr>', t, flags=re.IGNORECASE)
        lines = []
        for r in rows:
            cells = re.findall(r'<t[dh][^>]*>([\s\S]*?)</t[dh]>', r, flags=re.IGNORECASE)
            cells = [re.sub(r'<[^>]+>', '', c).strip() for c in cells]
            lines.append('| ' + ' | '.join(cells) + ' |')
        return chr(10).join(lines)
    for i, t in enumerate(tables):
        text = text.replace(f'@@TABLE{i}@@', _restore_table(t))
    return text.strip()

if len(sys.argv) < 2:
    print("Usage: import_jsonl.py <jsonl_path> [label]")
    sys.exit(1)

jsonl_path = sys.argv[1]
label = sys.argv[2] if len(sys.argv) > 2 else "data"

db_path = os.getenv("SEARCH_DB", "/root/search.db")
db = sqlite3.connect(db_path)
added = 0
errors = 0

with open(jsonl_path) as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        try:
            item = json.loads(line)
            title = (item.get("title") or "")[:500]
            page_url = (item.get("page_url") or "")[:1000]
            content = strip_html(item.get("content") or "")
            publish_date = (item.get("publish_date") or "")[:20]
            site_name = (item.get("site_name") or "unknown")[:100]

            # Attachments as JSON string
            att = item.get("attachments") or []
            attachments_str = json.dumps(att, ensure_ascii=False) if att else ""

            # Delete old FTS entry if this page_url already exists (rowid may change after REPLACE)
            db.execute("DELETE FROM gov_search WHERE rowid IN (SELECT rowid FROM gov_raw WHERE page_url = ?)", (page_url,))

            # Insert or replace raw data
            db.execute(
                "INSERT OR REPLACE INTO gov_raw (title, page_url, content, publish_date, site_name, source_url, status, attachments, script_name) VALUES (?, ?, ?, ?, ?, ?, 'synced', ?, 'import_jsonl_v2.py')",
                (title, page_url, content, publish_date, site_name, page_url, attachments_str)
            )

            # Get the new rowid and insert/update FTS incrementally
            new_rowid = db.execute("SELECT last_insert_rowid()").fetchone()[0]
            summary = content[:500] if content else title[:500]
            db.execute(
                "INSERT OR REPLACE INTO gov_search(rowid, title, site_name, summary) VALUES (?, ?, ?, ?)",
                (new_rowid, title, site_name, summary)
            )

            added += 1
        except Exception as e:
            errors += 1

db.commit()

total = db.execute("SELECT COUNT(*) FROM gov_raw").fetchone()[0]
db.close()

os.remove(jsonl_path)
print(f"\u2705 {label}: +{added} err{errors} total={total}")