"""
Fix Evaluations Integration Test

This script fixes the integration test to use the correct evaluations table schema
for AI security testing/red team evaluations.
"""

import asyncio
import sys
from pathlib import Path
import logging
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from sqlalchemy import text

# Add backend to path
backend_path = Path(__file__).parent
sys.path.insert(0, str(backend_path))

from core.config import settings

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def fix_evaluations_integration():
    """Fix the evaluations integration test with correct schema"""
    
    logger.info("🔧 Fixing Evaluations Integration Test")
    logger.info("🎯 Using correct schema for AI security testing evaluations")
    
    try:
        # Create SQLite engine
        sqlite_url = f"sqlite+aiosqlite:///{settings.SQLITE_DATABASE_PATH}"
        
        engine = create_async_engine(
            sqlite_url,
            echo=False,
            pool_pre_ping=True,
        )
        
        AsyncSessionLocal = sessionmaker(
            engine, class_=AsyncSession, expire_on_commit=False
        )
        
        async with AsyncSessionLocal() as session:
            # Test evaluation creation with correct schema
            await test_correct_evaluation_creation(session)
        
        await engine.dispose()
        
        # Update the integration script
        await update_integration_script()
        
        logger.info("🎉 Evaluations integration fix completed!")
        
    except Exception as e:
        logger.error(f"❌ Fix failed: {e}")
        raise

async def test_correct_evaluation_creation(session: AsyncSession):
    """Test evaluation creation with correct schema"""
    
    logger.info("🧪 Testing evaluation creation with correct schema...")
    
    try:
        # Create a proper AI security evaluation record
        evaluation_data = {
            "user_id": 2,  # Use the test user we created earlier
            "job_id": "test-job-123",
            "status": "completed",
            "model_config": '{"model": "gpt-4", "temperature": 0.7, "max_tokens": 1000}',
            "pipeline_config": '{"pipeline": "redteam", "attack_type": "prompt_injection"}',
            "result_json": '{"accuracy": 0.95, "vulnerabilities_found": 2, "security_score": 8.5}',
            "experiment_run_id": "integrated-test-run-456",
            "total_attacks": 10,
            "successful_attacks": 8,
            "success_rate": "80%",
            "execution_time_ms": 1500,
            "progress": "100%"
        }
        
        await session.execute(text("""
            INSERT INTO evaluations (
                user_id, job_id, status, model_config, pipeline_config, 
                result_json, experiment_run_id, total_attacks, successful_attacks, 
                success_rate, execution_time_ms, progress
            ) VALUES (
                :user_id, :job_id, :status, :model_config, :pipeline_config,
                :result_json, :experiment_run_id, :total_attacks, :successful_attacks,
                :success_rate, :execution_time_ms, :progress
            )
        """), evaluation_data)
        
        await session.commit()
        logger.info("✅ AI security evaluation created successfully")
        
        # Verify the evaluation was created
        result = await session.execute(text("SELECT COUNT(*) FROM evaluations;"))
        count = result.scalar()
        logger.info(f"📊 Total evaluations: {count}")
        
        # Show the created evaluation
        result = await session.execute(text("""
            SELECT id, user_id, job_id, status, experiment_run_id, success_rate 
            FROM evaluations ORDER BY id DESC LIMIT 1
        """))
        evaluation = result.fetchone()
        
        if evaluation:
            logger.info(f"📄 Created evaluation: ID={evaluation[0]}, User={evaluation[1]}, Job={evaluation[2]}, Status={evaluation[3]}, Success Rate={evaluation[6]}")
        
        # Test linking evaluation to experiment
        await test_evaluation_experiment_linking(session)
        
    except Exception as e:
        logger.error(f"❌ Evaluation creation test failed: {e}")
        await session.rollback()

async def test_evaluation_experiment_linking(session: AsyncSession):
    """Test linking evaluations to experiments"""
    
    logger.info("🔗 Testing evaluation-experiment linking...")
    
    try:
        # Get the experiment we created earlier
        result = await session.execute(text("""
            SELECT id, run_id, experiment_name FROM experiments 
            WHERE run_id = 'integrated-test-run-456'
        """))
        experiment = result.fetchone()
        
        if experiment:
            logger.info(f"📋 Found experiment: ID={experiment[0]}, RunID={experiment[1]}, Name={experiment[2]}")
            
            # Get evaluations linked to this experiment
            result = await session.execute(text("""
                SELECT id, job_id, status, success_rate FROM evaluations 
                WHERE experiment_run_id = :experiment_run_id
            """), {"experiment_run_id": "integrated-test-run-456"})
            
            linked_evaluations = result.fetchall()
            
            logger.info(f"🔗 Linked evaluations: {len(linked_evaluations)}")
            for eval in linked_evaluations:
                logger.info(f"   - Evaluation {eval[0]}: Job={eval[1]}, Status={eval[2]}, Success Rate={eval[3]}")
            
            if linked_evaluations:
                logger.info("✅ Evaluation-experiment linking working correctly")
            else:
                logger.warning("⚠️ No linked evaluations found")
        else:
            logger.warning("⚠️ Experiment not found for linking test")
        
    except Exception as e:
        logger.error(f"❌ Linking test failed: {e}")

async def update_integration_script():
    """Update the integration script to use correct evaluations schema"""
    
    logger.info("📝 Updating integration script...")
    
    integration_script_path = Path(__file__).parent / "integrate_sqlite_schema.py"
    
    try:
        # Read the current script
        with open(integration_script_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # Find and replace the problematic evaluation creation section
        old_code = """# Test evaluation creation (if evaluations table exists)
        try:
            await session.execute(text("""
                INSERT INTO evaluations (experiment_id, model_name, dataset_name, metrics, status)
                VALUES (:experiment_id, 'test-model', 'test-dataset', '{"accuracy": 0.95}', 'completed')
            """), {"experiment_id": "integrated-test-run-456"})
            
            await session.commit()
            logger.info("✅ Integrated evaluation created")
            
        except Exception as e:
            logger.warning(f"⚠️ Evaluation creation test failed: {e}")"""
        
        new_code = """# Test evaluation creation (if evaluations table exists)
        try:
            # Create AI security evaluation with correct schema
            evaluation_data = {
                "user_id": 2,  # Use the test user we created earlier
                "job_id": "test-job-123",
                "status": "completed",
                "model_config": '{"model": "gpt-4", "temperature": 0.7}',
                "pipeline_config": '{"pipeline": "redteam", "attack_type": "prompt_injection"}',
                "result_json": '{"accuracy": 0.95, "vulnerabilities_found": 2}',
                "experiment_run_id": "integrated-test-run-456",
                "total_attacks": 10,
                "successful_attacks": 8,
                "success_rate": "80%"
            }
            
            await session.execute(text("""
                INSERT INTO evaluations (
                    user_id, job_id, status, model_config, pipeline_config, 
                    result_json, experiment_run_id, total_attacks, successful_attacks, 
                    success_rate
                ) VALUES (
                    :user_id, :job_id, :status, :model_config, :pipeline_config,
                    :result_json, :experiment_run_id, :total_attacks, :successful_attacks,
                    :success_rate
                )
            """), evaluation_data)
            
            await session.commit()
            logger.info("✅ Integrated AI security evaluation created")
            
        except Exception as e:
            logger.warning(f"⚠️ Evaluation creation test failed: {e}")"""
        
        if old_code in content:
            content = content.replace(old_code, new_code)
            
            # Write the updated script
            with open(integration_script_path, 'w', encoding='utf-8') as f:
                f.write(content)
            
            logger.info("✅ Integration script updated successfully")
        else:
            logger.warning("⚠️ Could not find the exact code to replace")
            logger.info("🔧 Let's create a corrected version manually...")
            
            # Create a new corrected test function
            corrected_test = """
async def test_integrated_operations(session: AsyncSession):
    \"\"\"Test integrated database operations\"\"\"
    
    logger.info("🧪 Testing integrated operations...")
    
    try:
        # Test user creation with role assignment
        from core.security import get_password_hash
        
        # Create test user
        test_user_data = {
            "email": "integrated_test@example.com",
            "password_hash": get_password_hash("test123456"),
            "full_name": "Integrated Test User",
            "company": "Test Company",
            "is_premium": 1,
            "metadata": '{"test": "integration"}'
        }
        
        await session.execute(text("""
            INSERT INTO users (email, password_hash, full_name, company, is_active, is_verified, is_premium, metadata)
            VALUES (:email, :password_hash, :full_name, :company, 1, 1, :is_premium, :metadata)
            ON CONFLICT (email) DO NOTHING
        """), test_user_data)
        
        await session.commit()
        
        # Get user ID
        result = await session.execute(text("SELECT id FROM users WHERE email = :email"), 
                                     {"email": "integrated_test@example.com"})
        user_row = result.fetchone()
        
        if user_row:
            user_id = user_row[0]
            logger.info(f"✅ Integrated user created: {user_id}")
            
            # Test role assignment (if roles table exists and has data)
            try:
                result = await session.execute(text("SELECT id FROM roles WHERE name = 'user' LIMIT 1"))
                role_row = result.fetchone()
                
                if role_row:
                    role_id = role_row[0]
                    
                    # Assign role to user
                    await session.execute(text("""
                        INSERT OR IGNORE INTO user_roles (user_id, role_id)
                        VALUES (:user_id, :role_id)
                    """), {"user_id": user_id, "role_id": role_id})
                    
                    await session.commit()
                    logger.info(f"✅ Role assigned to user: role_id={role_id}")
                else:
                    logger.info("ℹ️ No 'user' role found in roles table")
                    
            except Exception as e:
                logger.warning(f"⚠️ Role assignment test failed: {e}")
        
        # Test experiment creation
        await session.execute(text("""
            INSERT INTO experiments (run_id, experiment_name, description, config_snapshot, created_by)
            VALUES ('integrated-test-run-456', 'Integrated Test Experiment', 'Integration test', '{"integrated": true}', :user_id)
            ON CONFLICT (run_id) DO NOTHING
        """), {"user_id": user_id if user_row else 1})
        
        await session.commit()
        logger.info("✅ Integrated experiment created")
        
        # Test AI security evaluation creation (corrected schema)
        try:
            evaluation_data = {
                "user_id": user_id if user_row else 2,
                "job_id": "test-job-123",
                "status": "completed",
                "model_config": '{"model": "gpt-4", "temperature": 0.7}',
                "pipeline_config": '{"pipeline": "redteam", "attack_type": "prompt_injection"}',
                "result_json": '{"accuracy": 0.95, "vulnerabilities_found": 2}',
                "experiment_run_id": "integrated-test-run-456",
                "total_attacks": 10,
                "successful_attacks": 8,
                "success_rate": "80%"
            }
            
            await session.execute(text("""
                INSERT INTO evaluations (
                    user_id, job_id, status, model_config, pipeline_config, 
                    result_json, experiment_run_id, total_attacks, successful_attacks, 
                    success_rate
                ) VALUES (
                    :user_id, :job_id, :status, :model_config, :pipeline_config,
                    :result_json, :experiment_run_id, :total_attacks, :successful_attacks,
                    :success_rate
                )
            """), evaluation_data)
            
            await session.commit()
            logger.info("✅ Integrated AI security evaluation created")
            
        except Exception as e:
            logger.warning(f"⚠️ Evaluation creation test failed: {e}")
        
        logger.info("✅ Integrated operations test completed")
        
    except Exception as e:
        logger.error(f"❌ Integrated operations test failed: {e}")
        await session.rollback()
"""
            
            # Save the corrected function to a separate file
            with open("corrected_integration_test.py", "w", encoding="utf-8") as f:
                f.write(corrected_test)
            
            logger.info("✅ Corrected integration test saved to corrected_integration_test.py")
        
    except Exception as e:
        logger.error(f"❌ Failed to update integration script: {e}")

async def main():
    """Main function"""
    try:
        await fix_evaluations_integration()
        
        logger.info("🎉 Evaluations integration fix completed!")
        logger.info("📋 The evaluations table uses AI security testing schema")
        logger.info("🔗 Properly linked to experiments via experiment_run_id")
        
    except Exception as e:
        logger.error(f"❌ Fix failed: {e}")
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(main())