#!/usr/bin/env python3
"""Migrate [表格] + HTML <table> to Markdown tables in search.db"""
import subprocess, sys, re
import html as html_mod

DB_PATH = '/mnt/data/search.db'

def html_table_to_md(html_table):
    """Parse HTML table and output Markdown table"""
    rows = re.findall(r'<tr[^>]*>(.*?)</tr>', html_table, re.DOTALL)
    md_rows = []
    first_row = True
    for row in rows:
        cells = re.findall(r'<t[dh][^>]*>(.*?)</t[dh]>', row, re.DOTALL)
        cell_texts = []
        for c in cells:
            ct = re.sub(r'<[^>]+>', '', c)
            ct = html_mod.unescape(ct)
            ct = re.sub(r'\s+', ' ', ct).strip()
            cell_texts.append(ct)
        if not any(cell_texts):
            continue
        md_line = ' | '.join(cell_texts)
        if first_row:
            md_rows.append(md_line)
            sep = ' | '.join(['---'] * len(cell_texts))
            md_rows.append(sep)
            first_row = False
        else:
            md_rows.append(md_line)
    return '\n'.join(md_rows)

def convert_content(content):
    """Replace [表格]\n<table>...</table> with Markdown table"""
    # Pattern: [表格]\n followed by HTML <table>...</table>
    pattern = r'\[表格\]\s*(<table[^>]*>.*?</table>)'
    def repl(m):
        md_table = html_table_to_md(m.group(1))
        return '\n\n' + md_table + '\n\n'
    
    new_content = re.sub(pattern, repl, content, flags=re.DOTALL)
    return new_content

def main():
    proc = subprocess.run(
        ["sqlite3", DB_PATH, "SELECT rowid FROM gov_raw WHERE content LIKE '%[表格]%'"],
        capture_output=True, text=True, timeout=120
    )
    if proc.returncode != 0:
        print("ERROR:", proc.stderr); sys.exit(1)
    
    rowids = [r.strip() for r in proc.stdout.strip().split('\n') if r.strip()]
    print("Found %d records" % len(rowids))
    
    updated = errors = 0
    for i, rowid in enumerate(rowids):
        r = subprocess.run(
            ["sqlite3", DB_PATH, "SELECT content FROM gov_raw WHERE rowid = %s" % rowid],
            capture_output=True, text=True, timeout=30
        )
        if r.returncode != 0:
            errors += 1; continue
        if not r.stdout.strip():
            continue
        
        new_c = convert_content(r.stdout)
        if new_c == r.stdout:
            continue
        
        escaped = new_c.replace("'", "''")
        u = subprocess.run(
            ["sqlite3", DB_PATH, "UPDATE gov_raw SET content = '%s' WHERE rowid = %s" % (escaped, rowid)],
            capture_output=True, text=True, timeout=30
        )
        if u.returncode != 0:
            print("ERROR update rowid %s: %s" % (rowid, u.stderr))
            errors += 1
        else:
            updated += 1
        
        if (i+1) % 50 == 0:
            print("  Progress: %d/%d" % (i+1, len(rowids)))
    
    print("\nDONE: %d updated, %d errors" % (updated, errors))

if __name__ == '__main__':
    main()