"""Data migration script for v2 schema migration.

Updates existing transcript records with appropriate v2 field values
and ensures backward compatibility with existing v1 data.
"""

import logging
from datetime import datetime, timezone
from typing import List, Dict, Any
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import SQLAlchemyError

from src.database.models import TranscriptionResult
from src.compatibility.backward_compatibility import TranscriptBackwardCompatibility

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


def migrate_existing_data(db_url: str) -> Dict[str, Any]:
    """Migrate existing transcript data to v2 format.
    
    Args:
        db_url: Database connection URL
        
    Returns:
        Dictionary with migration statistics
    """
    engine = create_engine(db_url)
    Session = sessionmaker(bind=engine)
    session = Session()
    
    migration_stats = {
        'total_transcripts': 0,
        'migrated_transcripts': 0,
        'skipped_transcripts': 0,
        'failed_transcripts': 0,
        'errors': []
    }
    
    try:
        # Get count of transcripts to migrate
        count = session.execute(text("SELECT COUNT(*) FROM transcription_results")).scalar()
        logger.info(f"Found {count} transcripts to migrate")
        migration_stats['total_transcripts'] = count
        
        # Get all transcripts
        transcripts = session.query(TranscriptionResult).all()
        
        for transcript in transcripts:
            try:
                # Skip if already v2
                if transcript.pipeline_version == 'v2':
                    migration_stats['skipped_transcripts'] += 1
                    logger.debug(f"Skipping transcript {transcript.id} - already v2")
                    continue
                
                # Migrate transcript to v2
                TranscriptBackwardCompatibility.migrate_v1_to_v2(transcript)
                
                # Update timestamp
                transcript.updated_at = datetime.now(timezone.utc)
                
                migration_stats['migrated_transcripts'] += 1
                logger.debug(f"Migrated transcript {transcript.id} to v2")
                
            except Exception as e:
                migration_stats['failed_transcripts'] += 1
                error_msg = f"Failed to migrate transcript {transcript.id}: {str(e)}"
                migration_stats['errors'].append(error_msg)
                logger.error(error_msg)
        
        # Commit all changes
        session.commit()
        logger.info("Migration completed successfully")
        
    except SQLAlchemyError as e:
        session.rollback()
        error_msg = f"Database error during migration: {str(e)}"
        migration_stats['errors'].append(error_msg)
        logger.error(error_msg)
        raise
    except Exception as e:
        session.rollback()
        error_msg = f"Unexpected error during migration: {str(e)}"
        migration_stats['errors'].append(error_msg)
        logger.error(error_msg)
        raise
    finally:
        session.close()
    
    # Log final statistics
    logger.info(f"Migration completed: {migration_stats['migrated_transcripts']} migrated, "
                f"{migration_stats['skipped_transcripts']} skipped, "
                f"{migration_stats['failed_transcripts']} failed")
    
    return migration_stats


def migrate_specific_transcripts(db_url: str, transcript_ids: List[str]) -> Dict[str, Any]:
    """Migrate specific transcripts to v2 format.
    
    Args:
        db_url: Database connection URL
        transcript_ids: List of transcript IDs to migrate
        
    Returns:
        Dictionary with migration statistics
    """
    engine = create_engine(db_url)
    Session = sessionmaker(bind=engine)
    session = Session()
    
    migration_stats = {
        'total_transcripts': len(transcript_ids),
        'migrated_transcripts': 0,
        'skipped_transcripts': 0,
        'failed_transcripts': 0,
        'errors': []
    }
    
    try:
        for transcript_id in transcript_ids:
            try:
                # Get transcript
                transcript = session.query(TranscriptionResult).filter(
                    TranscriptionResult.id == transcript_id
                ).first()
                
                if not transcript:
                    error_msg = f"Transcript {transcript_id} not found"
                    migration_stats['errors'].append(error_msg)
                    migration_stats['failed_transcripts'] += 1
                    logger.error(error_msg)
                    continue
                
                # Skip if already v2
                if transcript.pipeline_version == 'v2':
                    migration_stats['skipped_transcripts'] += 1
                    logger.debug(f"Skipping transcript {transcript_id} - already v2")
                    continue
                
                # Migrate transcript to v2
                TranscriptBackwardCompatibility.migrate_v1_to_v2(transcript)
                
                # Update timestamp
                transcript.updated_at = datetime.now(timezone.utc)
                
                migration_stats['migrated_transcripts'] += 1
                logger.info(f"Migrated transcript {transcript_id} to v2")
                
            except Exception as e:
                migration_stats['failed_transcripts'] += 1
                error_msg = f"Failed to migrate transcript {transcript_id}: {str(e)}"
                migration_stats['errors'].append(error_msg)
                logger.error(error_msg)
        
        # Commit all changes
        session.commit()
        logger.info("Specific migration completed successfully")
        
    except SQLAlchemyError as e:
        session.rollback()
        error_msg = f"Database error during migration: {str(e)}"
        migration_stats['errors'].append(error_msg)
        logger.error(error_msg)
        raise
    except Exception as e:
        session.rollback()
        error_msg = f"Unexpected error during migration: {str(e)}"
        migration_stats['errors'].append(error_msg)
        logger.error(error_msg)
        raise
    finally:
        session.close()
    
    return migration_stats


def validate_migration(db_url: str) -> Dict[str, Any]:
    """Validate the migration results.
    
    Args:
        db_url: Database connection URL
        
    Returns:
        Dictionary with validation results
    """
    engine = create_engine(db_url)
    Session = sessionmaker(bind=engine)
    session = Session()
    
    validation_results = {
        'total_transcripts': 0,
        'v1_transcripts': 0,
        'v2_transcripts': 0,
        'v2_with_features': 0,
        'validation_errors': []
    }
    
    try:
        # Get all transcripts
        transcripts = session.query(TranscriptionResult).all()
        validation_results['total_transcripts'] = len(transcripts)
        
        for transcript in transcripts:
            try:
                if transcript.pipeline_version == 'v1':
                    validation_results['v1_transcripts'] += 1
                elif transcript.pipeline_version == 'v2':
                    validation_results['v2_transcripts'] += 1
                    
                    # Check if v2 transcript has v2 features
                    if TranscriptBackwardCompatibility.has_v2_features(transcript):
                        validation_results['v2_with_features'] += 1
                
            except Exception as e:
                error_msg = f"Validation error for transcript {transcript.id}: {str(e)}"
                validation_results['validation_errors'].append(error_msg)
                logger.error(error_msg)
        
        # Log validation results
        logger.info(f"Validation completed: {validation_results['total_transcripts']} total, "
                    f"{validation_results['v1_transcripts']} v1, "
                    f"{validation_results['v2_transcripts']} v2, "
                    f"{validation_results['v2_with_features']} v2 with features")
        
    except Exception as e:
        error_msg = f"Validation failed: {str(e)}"
        validation_results['validation_errors'].append(error_msg)
        logger.error(error_msg)
        raise
    finally:
        session.close()
    
    return validation_results


def rollback_migration(db_url: str) -> Dict[str, Any]:
    """Rollback v2 migration by setting pipeline_version back to v1.
    
    Args:
        db_url: Database connection URL
        
    Returns:
        Dictionary with rollback statistics
    """
    engine = create_engine(db_url)
    Session = sessionmaker(bind=engine)
    session = Session()
    
    rollback_stats = {
        'total_transcripts': 0,
        'rolled_back_transcripts': 0,
        'skipped_transcripts': 0,
        'failed_transcripts': 0,
        'errors': []
    }
    
    try:
        # Get count of v2 transcripts
        count = session.execute(text(
            "SELECT COUNT(*) FROM transcription_results WHERE pipeline_version = 'v2'"
        )).scalar()
        logger.info(f"Found {count} v2 transcripts to rollback")
        rollback_stats['total_transcripts'] = count
        
        # Get all v2 transcripts
        v2_transcripts = session.query(TranscriptionResult).filter(
            TranscriptionResult.pipeline_version == 'v2'
        ).all()
        
        for transcript in v2_transcripts:
            try:
                # Set pipeline version back to v1
                transcript.pipeline_version = 'v1'
                
                # Clear v2-specific fields
                transcript.enhanced_content = None
                transcript.diarization_content = None
                transcript.merged_content = None
                transcript.domain_used = None
                transcript.accuracy_estimate = None
                transcript.speaker_count = None
                transcript.quality_warnings = None
                transcript.processing_metadata = None
                
                # Update timestamp
                transcript.updated_at = datetime.now(timezone.utc)
                
                rollback_stats['rolled_back_transcripts'] += 1
                logger.debug(f"Rolled back transcript {transcript.id} to v1")
                
            except Exception as e:
                rollback_stats['failed_transcripts'] += 1
                error_msg = f"Failed to rollback transcript {transcript.id}: {str(e)}"
                rollback_stats['errors'].append(error_msg)
                logger.error(error_msg)
        
        # Commit all changes
        session.commit()
        logger.info("Rollback completed successfully")
        
    except SQLAlchemyError as e:
        session.rollback()
        error_msg = f"Database error during rollback: {str(e)}"
        rollback_stats['errors'].append(error_msg)
        logger.error(error_msg)
        raise
    except Exception as e:
        session.rollback()
        error_msg = f"Unexpected error during rollback: {str(e)}"
        rollback_stats['errors'].append(error_msg)
        logger.error(error_msg)
        raise
    finally:
        session.close()
    
    # Log final statistics
    logger.info(f"Rollback completed: {rollback_stats['rolled_back_transcripts']} rolled back, "
                f"{rollback_stats['failed_transcripts']} failed")
    
    return rollback_stats


if __name__ == "__main__":
    """Run migration as a script."""
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python data_migration.py <database_url> [action]")
        print("Actions: migrate, validate, rollback")
        sys.exit(1)
    
    db_url = sys.argv[1]
    action = sys.argv[2] if len(sys.argv) > 2 else "migrate"
    
    try:
        if action == "migrate":
            stats = migrate_existing_data(db_url)
            print(f"Migration completed: {stats}")
        elif action == "validate":
            results = validate_migration(db_url)
            print(f"Validation completed: {results}")
        elif action == "rollback":
            stats = rollback_migration(db_url)
            print(f"Rollback completed: {stats}")
        else:
            print(f"Unknown action: {action}")
            sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
