"""Fix gov_raw table corruption by rebuilding the table"""
import sqlite3, sys

DB = '/root/search.db.fixed'

print(f"Opening {DB}...")
db = sqlite3.connect(DB)
db.execute("PRAGMA journal_mode=WAL")
db.execute("PRAGMA busy_timeout=30000")

# Get gov_raw schema
c = db.cursor()
c.execute("SELECT sql FROM sqlite_master WHERE name='gov_raw' AND type='table'")
row = c.fetchone()
schema_sql = row[0]
print(f"Current schema: {schema_sql[:100]}...")

# Create new table
new_schema = schema_sql.replace('"gov_raw"', '"gov_raw_new"').replace('gov_raw', 'gov_raw_new')
print(f"\nCreating new table...")
db.execute(f"DROP TABLE IF EXISTS gov_raw_new")
db.execute(new_schema)
db.commit()
print("New table created")

# Copy data in batches
c.execute("SELECT COUNT(*) FROM gov_raw")
total = c.fetchone()[0]
print(f"Copying {total} rows...")

BATCH = 2000
offset = 0
while offset < total:
    db.execute(f"INSERT INTO gov_raw_new SELECT * FROM gov_raw LIMIT {BATCH} OFFSET {offset}")
    db.commit()
    offset += BATCH
    if offset % 10000 == 0:
        print(f"  {min(offset, total)}/{total}")

print(f"  {total}/{total}")

# Get indices for gov_raw
c.execute("SELECT sql FROM sqlite_master WHERE type='index' AND tbl_name='gov_raw' AND sql IS NOT NULL")
indices = c.fetchall()

# Drop old table
print("Dropping old table...")
db.execute("DROP TABLE IF EXISTS gov_raw_old")
db.execute("ALTER TABLE gov_raw RENAME TO gov_raw_old")
db.commit()

# Rename new table to gov_raw
print("Renaming new table...")
db.execute("ALTER TABLE gov_raw_new RENAME TO gov_raw")
db.commit()

# Recreate indices
print("Recreating indices...")
for idx in indices:
    idx_sql = idx[0].replace('"gov_raw"', '"gov_raw"')
    try:
        db.execute(idx_sql)
    except Exception as e:
        print(f"  Index error: {e}")
db.commit()

# Verify
cnt = db.execute("SELECT COUNT(*) FROM gov_raw").fetchone()[0]
print(f"\nNew gov_raw count: {cnt}")

# Test write
try:
    db.execute("INSERT INTO gov_raw(site_name,title,page_url,publish_date) VALUES('test_ok','test','test','2024-01-01')")
    db.commit()
    print("Write: OK ✓")
    db.execute("DELETE FROM gov_raw WHERE site_name='test_ok'")
    db.commit()
    print("Delete: OK ✓")
except Exception as e:
    print(f"Write/Delete FAILED: {e}")

# Drop old table
print("Cleaning up old table...")
db.execute("DROP TABLE IF EXISTS gov_raw_old")
db.commit()

# Rebuild FTS
print("Rebuilding FTS...")
try:
    db.execute("INSERT INTO gov_search(gov_search) VALUES('rebuild')")
    db.commit()
    print("FTS rebuild: OK ✓")
except Exception as e:
    print(f"FTS rebuild: {e}")

db.close()
print("\n✅ All done!")