#!/usr/bin/env python3
"""DB Recovery: merge corrupted DB .dumps into recovered base DB.

Strategy:
1. Start from mofin.db.recovered (clean, all tables, integrity OK)
2. For 'dead' tables (not directly readable), import from .dump files
3. For 'alive' tables, compare row counts and keep the fuller version
4. For holding_strategies, merge dump data (71 rows) + recovered (22 rows)

Usage: python3 scripts/db_recovery.py
"""

import sqlite3
import os
import sys
import shutil
from datetime import datetime

RECOVERED = "/home/hmo/MoFin/data/mofin.db.recovered"
OUTPUT = "/home/hmo/MoFin/data/mofin.db.new_recovered"
CORRUPTED = "/home/hmo/MoFin/data/mofin.db"

DUMP_DIR = "/tmp"

# Which tables are DEAD in corrupted DB (have page corruption)
DEAD_TABLES = [
    "price_events",        # 2700 rows in dump
    "stock_weekly",        # 646 rows in dump
    "stock_monthly",       # 517 rows in dump
    "stock_daily",         # 0 rows in dump (empty)
    "market_snapshots",    # 808 rows in dump
    "sector_snapshots",    # 62087 rows in dump
    "macro_raw_news",      # 15499 rows in dump
    "strategy_history",    # 12 rows in dump
    "health_check_log",    # 20 rows in dump
]

# Tables that are ALIVE in corrupted + have .dump available
# For these we compare row counts and keep the fuller version
ALIVE_WITH_DUMP = {
    "holding_strategies": "/tmp/strategies_dump.sql",  # 71 in dump vs 22 in recovered
    "candidates": None,                                # 216 in corrupted vs 139 in recovered
}

def get_row_counts(conn):
    """Get row counts for all tables."""
    c = conn.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
    counts = {}
    for r in c.fetchall():
        tname = r[0]
        try:
            cnt = conn.execute(f'SELECT COUNT(*) FROM "{tname}"').fetchone()[0]
            counts[tname] = cnt
        except Exception as e:
            counts[tname] = -1
    return counts

def import_dump(conn, dump_path, table_name):
    """Import SQL dump into database, clearing table first."""
    if not os.path.exists(dump_path):
        print(f"  DUMP FILE NOT FOUND: {dump_path}")
        return 0
    
    # Read the dump file
    with open(dump_path, 'r') as f:
        sql = f.read()
    
    # Extract CREATE TABLE and INSERT statements
    # First, clear existing data
    conn.execute(f'DELETE FROM "{table_name}"')
    conn.commit()
    
    # Execute all INSERT statements from the dump
    inserts_count = 0
    cur = conn.cursor()
    
    for line in sql.split('\n'):
        line = line.strip()
        if line.startswith('INSERT INTO'):
            try:
                cur.execute(line)
                inserts_count += 1
            except Exception as e:
                print(f"  ERROR on insert: {str(e)[:100]}")
                conn.rollback()
                return inserts_count
    
    conn.commit()
    return inserts_count

def merge_strategies(conn, dump_path):
    """Special handling for holding_strategies:
    - Import from dump (71 rows)
    - Supplement with recovered rows that are NOT already in dump
    """
    if not os.path.exists(dump_path):
        print("  Dump file not found, keeping recovered strategies")
        return
    
    # Read dump SQL
    with open(dump_path, 'r') as f:
        dump_sql = f.read()
    
    # First, collect all IDs already in the recovered DB
    existing_ids = set()
    try:
        c = conn.execute('SELECT id FROM holding_strategies')
        existing_ids = {r[0] for r in c.fetchall()}
    except:
        pass
    
    # Check if table still has data
    c = conn.execute('SELECT COUNT(*) FROM holding_strategies')
    recovered_count = c.fetchone()[0]
    print(f"  Recovered has {recovered_count} strategies")
    
    # Parse dump to get IDs
    dump_ids = set()
    for line in dump_sql.split('\n'):
        line = line.strip()
        if line.startswith('INSERT INTO holding_strategies VALUES('):
            # Extract the first value (id)
            try:
                # Find the first value after VALUES(
                vals_start = line.index('VALUES(') + 7
                # First value is the id, which may be quoted or not
                rest = line[vals_start:]
                if rest.startswith("'"):
                    # String id - unlikely but handle
                    end = rest.index("'", 1)
                    val = rest[1:end]
                else:
                    # Numeric id
                    end = rest.index(',')
                    val = rest[:end].strip()
                if val:
                    dump_ids.add(int(val) if val.isdigit() else val)
            except Exception as e:
                pass
    
    print(f"  Dump has {len(dump_ids)} unique strategy IDs")
    
    # Delete all existing strategies from the table
    conn.execute('DELETE FROM holding_strategies')
    conn.commit()
    
    # Now import all from dump
    dump_inserted = 0
    cur = conn.cursor()
    for line in dump_sql.split('\n'):
        line = line.strip()
        if line.startswith('INSERT INTO holding_strategies VALUES('):
            try:
                cur.execute(line)
                dump_inserted += 1
            except Exception as e:
                pass
    conn.commit()
    print(f"  Imported {dump_inserted} strategies from dump")
    
    return dump_inserted

def main():
    print(f"=== MoFin DB Recovery at {datetime.now()} ===")
    
    # Verify inputs
    if not os.path.exists(RECOVERED):
        print(f"ERROR: Recovered DB not found at {RECOVERED}")
        sys.exit(1)
    
    # Step 1: Copy recovered DB as base
    print(f"\nStep 1: Copying {RECOVERED} -> {OUTPUT}")
    shutil.copy2(RECOVERED, OUTPUT)
    
    # Step 2: Connect and verify integrity
    conn = sqlite3.connect(OUTPUT)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")
    
    c = conn.execute("PRAGMA quick_check")
    result = c.fetchone()[0]
    print(f"  Base integrity: {result}")
    
    if result != "ok":
        print("  ERROR: Base DB failed integrity check!")
        conn.close()
        sys.exit(1)
    
    # Step 3: Get baseline row counts
    before = get_row_counts(conn)
    print(f"\nStep 2: Row counts before merge:")
    for t, cnt in sorted(before.items(), key=lambda x: x[0]):
        print(f"  {t}: {cnt}")
    
    # Step 4: Import dead tables from .dump
    print(f"\nStep 3: Importing dead table dumps...")
    for table in DEAD_TABLES:
        dump_path = os.path.join(DUMP_DIR, f"{table}_dump.sql")
        if not os.path.exists(dump_path):
            print(f"  SKIP {table}: dump not found at {dump_path}")
            continue
        
        print(f"  Importing {table}...")
        count = import_dump(conn, dump_path, table)
        print(f"    Imported {count} rows into {table}")
    
    # Step 5: Handle holding_strategies specially
    print(f"\nStep 4: Merging holding_strategies...")
    dump_path = f"{DUMP_DIR}/strategies_dump.sql"
    merge_strategies(conn, dump_path)
    
    # Step 6: Verify final state
    print(f"\nStep 5: Final row counts:")
    after = get_row_counts(conn)
    for t in sorted(after.keys()):
        b = before.get(t, 0)
        a = after.get(t, 0)
        if b != a:
            print(f"  {t}: {b} -> {a} ({'+' if a > b else ''}{a-b})")
        else:
            print(f"  {t}: {a} (unchanged)")
    
    # Step 7: Final integrity check
    c = conn.execute("PRAGMA quick_check")
    result = c.fetchone()[0]
    print(f"\nFinal integrity: {result}")
    
    conn.close()
    
    print(f"\n=== Recovery complete. Output: {OUTPUT} ===")
    print(f"File size: {os.path.getsize(OUTPUT):,} bytes")
    
    # Provide instructions
    print(f"\nTo deploy:")
    print(f"  cd /home/hmo/MoFin/data")
    print(f"  cp mofin.db mofin.db.corrupted_backup.2")
    print(f"  cp mofin.db.new_recovered mofin.db")

if __name__ == "__main__":
    main()