#!/usr/bin/env python3
"""
Cleanup Test Data Script
Removes all test data from the database that was created during testing.
"""

import sqlite3
import os
import sys

# Database path
DB_PATH = os.path.join(os.path.dirname(__file__), '..', 'database', 'vera_medical.db')

def get_db_conn():
    """Get database connection with proper settings"""
    conn = sqlite3.connect(DB_PATH, timeout=30.0)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=30000")
    return conn

def cleanup_test_data():
    """Remove all test data from database"""
    print("=" * 60)
    print("CLEANUP TEST DATA")
    print("=" * 60)
    
    if not os.path.exists(DB_PATH):
        print(f"❌ Database not found at {DB_PATH}")
        sys.exit(1)
    
    conn = get_db_conn()
    cursor = conn.cursor()
    
    # Tables to clean up in dependency order (delete child records first)
    cleanup_tables = [
        # Child tables first
        ('bring_forward_tasks', 'task_title'),
        ('case_reports', 'file_name'),
        ('generated_letters', 'letter_subject'),
        ('invoice_line_items', 'description'),
        ('appointment_meds', 'file_name'),
        ('medical_documentation', 'file_name'),
        ('appointment_additional_services', 'service_type'),
        ('appointment_letters', 'file_name'),
        # Then parent tables
        ('invoices', 'invoice_number'),
        ('appointments', 'notes'),
        ('reports', 'report_title'),
        ('cases', 'case_notes'),
        ('claimants', 'first_name'),
        ('clients', 'company_name'),
        ('lawyer_contacts', 'first_name'),
        ('lawyer_firms', 'firm_name')
    ]
    
    total_deleted = 0
    
    for table, field in cleanup_tables:
        try:
            # Check if table exists
            cursor.execute(f"SELECT name FROM sqlite_master WHERE type='table' AND name='{table}'")
            if not cursor.fetchone():
                continue
            
            # Check if column exists
            cursor.execute(f"PRAGMA table_info({table})")
            columns = [row[1] for row in cursor.fetchall()]
            if field not in columns:
                continue
            
            # Count test records
            cursor.execute(f"""
                SELECT COUNT(*) FROM {table} 
                WHERE {field} LIKE 'TEST%' 
                   OR {field} LIKE '%TEST_PHASE%'
                   OR {field} LIKE '%TEST_AUTO%'
                   OR {field} LIKE 'test_%'
            """)
            count = cursor.fetchone()[0]
            
            if count > 0:
                # Delete test records
                cursor.execute(f"""
                    DELETE FROM {table} 
                    WHERE {field} LIKE 'TEST%' 
                       OR {field} LIKE '%TEST_PHASE%'
                       OR {field} LIKE '%TEST_AUTO%'
                       OR {field} LIKE 'test_%'
                """)
                deleted = cursor.rowcount
                total_deleted += deleted
                print(f"✅ {table}: Deleted {deleted} test record(s)")
        except Exception as e:
            print(f"⚠️  {table}: Error - {e}")
    
    # Also clean up by other patterns
    # Clean up cases with test notes
    try:
        cursor.execute("""
            DELETE FROM cases 
            WHERE case_notes LIKE '%TEST%' 
               OR service_location LIKE 'TEST%'
        """)
        if cursor.rowcount > 0:
            total_deleted += cursor.rowcount
            print(f"✅ cases: Deleted {cursor.rowcount} additional test record(s)")
    except Exception as e:
        print(f"⚠️  cases: Error - {e}")
    
    # Clean up claimants with test names
    try:
        cursor.execute("""
            DELETE FROM claimants 
            WHERE first_name LIKE 'TEST%' 
               OR last_name LIKE 'TEST%'
               OR first_name LIKE 'test_%'
               OR last_name LIKE 'test_%'
        """)
        if cursor.rowcount > 0:
            total_deleted += cursor.rowcount
            print(f"✅ claimants: Deleted {cursor.rowcount} additional test record(s)")
    except Exception as e:
        print(f"⚠️  claimants: Error - {e}")
    
    # Clean up clients with test company names
    try:
        cursor.execute("""
            DELETE FROM clients 
            WHERE company_name LIKE 'TEST%' 
               OR company_name LIKE 'test_%'
               OR client_name LIKE 'TEST%'
        """)
        if cursor.rowcount > 0:
            total_deleted += cursor.rowcount
            print(f"✅ clients: Deleted {cursor.rowcount} additional test record(s)")
    except Exception as e:
        print(f"⚠️  clients: Error - {e}")
    
    # Clean up appointments with test data
    try:
        cursor.execute("""
            DELETE FROM appointments 
            WHERE notes LIKE '%TEST%' 
               OR doctor_name LIKE '%TEST%'
               OR notes LIKE 'test_%'
        """)
        if cursor.rowcount > 0:
            total_deleted += cursor.rowcount
            print(f"✅ appointments: Deleted {cursor.rowcount} additional test record(s)")
    except Exception as e:
        print(f"⚠️  appointments: Error - {e}")
    
    # Clean up any remaining claimants
    try:
        cursor.execute("""
            DELETE FROM claimants 
            WHERE first_name LIKE '%TEST%' 
               OR last_name LIKE '%TEST%'
               OR first_name LIKE 'test_%'
               OR last_name LIKE 'test_%'
        """)
        if cursor.rowcount > 0:
            total_deleted += cursor.rowcount
            print(f"✅ claimants: Deleted {cursor.rowcount} additional test record(s)")
    except Exception as e:
        print(f"⚠️  claimants: Error - {e}")
    
    # Commit all changes
    conn.commit()
    conn.close()
    
    print("=" * 60)
    print(f"✅ Cleanup complete! Total records deleted: {total_deleted}")
    print("=" * 60)
    
    return total_deleted

if __name__ == '__main__':
    try:
        deleted = cleanup_test_data()
        if deleted == 0:
            print("ℹ️  No test data found in database")
        sys.exit(0)
    except Exception as e:
        print(f"❌ Error during cleanup: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)