#!/usr/bin/env python3
"""
Utility to detect and clear SQLite database locks
"""

import os
import sqlite3
import argparse
from pathlib import Path
import sys
sys.path.append('/Users/andrew/Desktop/cc/text2sql')

from cached_sql_executor import check_database_lock

def find_lock_files(db_root: str):
    """Find SQLite lock files (WAL, SHM, journal files)"""
    lock_files = []
    
    for root, dirs, files in os.walk(db_root):
        for file in files:
            if file.endswith(('.sqlite-wal', '.sqlite-shm', '.sqlite-journal')):
                lock_files.append(os.path.join(root, file))
    
    return lock_files

def clear_lock_files(lock_files: list, dry_run: bool = True):
    """Remove SQLite lock files"""
    if not lock_files:
        print("✅ No lock files found")
        return
    
    print(f"Found {len(lock_files)} lock files:")
    for lock_file in lock_files:
        print(f"  - {lock_file}")
    
    if dry_run:
        print("\n🔍 DRY RUN - Would remove these files (use --force to actually remove)")
        return
    
    print(f"\n🗑️  Removing {len(lock_files)} lock files...")
    removed = 0
    for lock_file in lock_files:
        try:
            os.remove(lock_file)
            print(f"  ✅ Removed: {lock_file}")
            removed += 1
        except Exception as e:
            print(f"  ❌ Failed to remove {lock_file}: {e}")
    
    print(f"\n✅ Successfully removed {removed}/{len(lock_files)} lock files")

def test_database_access(db_root: str):
    """Test access to all SQLite databases"""
    print("Testing database access...")
    
    db_files = []
    for root, dirs, files in os.walk(db_root):
        for file in files:
            if file.endswith('.sqlite'):
                db_files.append(os.path.join(root, file))
    
    if not db_files:
        print("❌ No .sqlite files found")
        return
    
    locked_dbs = []
    accessible_dbs = []
    
    for db_file in db_files:
        try:
            is_locked = check_database_lock(db_file)
            if is_locked:
                locked_dbs.append(db_file)
                print(f"  🔒 LOCKED: {db_file}")
            else:
                accessible_dbs.append(db_file)
                print(f"  ✅ OK: {db_file}")
        except Exception as e:
            locked_dbs.append(db_file)
            print(f"  ❌ ERROR: {db_file} - {e}")
    
    print(f"\n📊 Summary:")
    print(f"  Total databases: {len(db_files)}")
    print(f"  Accessible: {len(accessible_dbs)}")
    print(f"  Locked/Problem: {len(locked_dbs)}")
    
    return locked_dbs

def force_unlock_database(db_path: str):
    """Force unlock a database using various methods"""
    print(f"Attempting to force unlock: {db_path}")
    
    methods_tried = []
    
    # Method 1: Try to open with immediate timeout and close
    try:
        conn = sqlite3.connect(db_path, timeout=0.1)
        conn.execute("BEGIN IMMEDIATE;")
        conn.rollback()
        conn.close()
        methods_tried.append("immediate_transaction")
        print(f"  ✅ Method 1 (immediate transaction) succeeded")
        return True
    except Exception as e:
        methods_tried.append(f"immediate_transaction_failed: {e}")
        print(f"  ❌ Method 1 failed: {e}")
    
    # Method 2: Try WAL mode checkpoint
    try:
        conn = sqlite3.connect(db_path, timeout=0.1)
        conn.execute("PRAGMA wal_checkpoint(TRUNCATE);")
        conn.close()
        methods_tried.append("wal_checkpoint")
        print(f"  ✅ Method 2 (WAL checkpoint) succeeded")
        return True
    except Exception as e:
        methods_tried.append(f"wal_checkpoint_failed: {e}")
        print(f"  ❌ Method 2 failed: {e}")
    
    # Method 3: Remove associated lock files
    try:
        base_path = db_path.rsplit('.', 1)[0]
        lock_files = [
            f"{db_path}-wal",
            f"{db_path}-shm", 
            f"{db_path}-journal"
        ]
        
        removed_files = []
        for lock_file in lock_files:
            if os.path.exists(lock_file):
                os.remove(lock_file)
                removed_files.append(lock_file)
        
        if removed_files:
            methods_tried.append(f"removed_lock_files: {removed_files}")
            print(f"  ✅ Method 3 (remove lock files) removed: {removed_files}")
            return True
        else:
            methods_tried.append("no_lock_files_found")
            print(f"  ❌ Method 3: No lock files found to remove")
    except Exception as e:
        methods_tried.append(f"remove_lock_files_failed: {e}")
        print(f"  ❌ Method 3 failed: {e}")
    
    print(f"  ❌ All unlock methods failed for {db_path}")
    return False

def main():
    parser = argparse.ArgumentParser(description='Clear SQLite database locks')
    parser.add_argument('--db_root', default='benchmark_resources/datasets/dev/dev_20240627/dev_databases',
                       help='Root directory containing SQLite databases')
    parser.add_argument('--force', action='store_true', 
                       help='Actually remove lock files (default is dry run)')
    parser.add_argument('--test', action='store_true',
                       help='Test database access without removing files')
    parser.add_argument('--unlock', action='store_true',
                       help='Attempt to force unlock locked databases')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.db_root):
        print(f"❌ Database root not found: {args.db_root}")
        return 1
    
    print(f"🔍 Checking databases in: {args.db_root}")
    
    # Find and optionally remove lock files
    lock_files = find_lock_files(args.db_root)
    clear_lock_files(lock_files, dry_run=not args.force)
    
    # Test database access
    if args.test or args.unlock:
        locked_dbs = test_database_access(args.db_root)
        
        # Try to force unlock if requested
        if args.unlock and locked_dbs:
            print(f"\n🔓 Attempting to force unlock {len(locked_dbs)} databases...")
            unlocked = 0
            for db_path in locked_dbs:
                if force_unlock_database(db_path):
                    unlocked += 1
            
            print(f"\n📊 Unlock Summary:")
            print(f"  Attempted: {len(locked_dbs)}")
            print(f"  Successfully unlocked: {unlocked}")
            
            # Re-test after unlock attempts
            if unlocked > 0:
                print(f"\n🔄 Re-testing database access...")
                test_database_access(args.db_root)
    
    return 0

if __name__ == "__main__":
    exit(main())