"""Transcription commands for the Trax CLI."""

import asyncio
import json
import click
from pathlib import Path
from rich.progress import Progress, TextColumn, BarColumn, TimeElapsedColumn

from ..utils import console, get_media_extensions
from ...services.factories import create_transcription_service
from ..progress import CliProgressRenderer
from ...services.media_types import ProcessingProgress
from ...services.multi_pass_transcription import MultiPassTranscriptionPipeline


@click.command()
@click.argument('file', type=click.Path(exists=True))
@click.option('--json', 'output_format', flag_value='json', default=True, help='Output as JSON')
@click.option('--txt', 'output_format', flag_value='txt', help='Output as plain text')
@click.option('--v1', 'version', flag_value='v1', default=True, help='Use v1 pipeline (Whisper only)')
@click.option('--v2', 'version', flag_value='v2', help='Use v2 pipeline (Whisper + Enhancement)')
@click.option('--multi-pass', 'version', flag_value='multi-pass', help='Use multi-pass pipeline (fast pass + refinement + enhancement)')
@click.option('--min-accuracy', default=80.0, help='Minimum accuracy threshold (default: 80%)')
@click.option('--confidence-threshold', default=0.85, help='Confidence threshold for refinement (default: 0.85)')
@click.option('--domain', type=click.Choice(['general', 'technical', 'medical', 'academic']), help='Domain for enhancement pass')
@click.option('--diarize', is_flag=True, help='Enable speaker diarization in multi-pass mode')
@click.option('--m3-status', is_flag=True, help='Show M3 optimization status before transcription')
def transcribe(file, output_format, version, min_accuracy, confidence_threshold, domain, diarize, m3_status):
    """Transcribe a single audio file."""
    async def _transcribe():
        try:
            file_path = Path(file)
            
            if version == 'multi-pass':
                # Use multi-pass transcription pipeline
                pipeline = MultiPassTranscriptionPipeline()
                pipeline.confidence_threshold = confidence_threshold
                
                console.print(f"[cyan]Using multi-pass pipeline[/cyan]")
                console.print(f"[dim]• Confidence threshold: {confidence_threshold}[/dim]")
                if domain:
                    console.print(f"[dim]• Domain: {domain}[/dim]")
                if diarize:
                    console.print(f"[dim]• Speaker diarization: enabled[/dim]")
                
            else:
                # Initialize legacy services for v1/v2
                transcription_service = create_transcription_service()
                await transcription_service.initialize()
                
                # Configure transcription
                from ...services.local_transcription_service import LocalTranscriptionConfig
                config = LocalTranscriptionConfig(
                    model="distil-large-v3",
                    chunk_size_seconds=600
                )
                
                # Show M3 optimization status if requested
                if m3_status and hasattr(transcription_service, 'get_m3_optimization_status'):
                    try:
                        m3_status_info = transcription_service.get_m3_optimization_status()
                        console.print("\n[bold cyan]M3 Optimization Status:[/bold cyan]")
                        console.print(f"  Device: {m3_status_info.get('device', 'unknown')}")
                        console.print(f"  MPS Available: {m3_status_info.get('mps_available', False)}")
                        console.print(f"  M3 Preprocessing: {m3_status_info.get('m3_preprocessing_enabled', False)}")
                        console.print(f"  Hardware Acceleration: {m3_status_info.get('hardware_acceleration_enabled', False)}")
                        console.print(f"  VideoToolbox Support: {m3_status_info.get('videotoolbox_support', False)}")
                        console.print(f"  Compute Type: {m3_status_info.get('compute_type', 'unknown')}")
                        console.print()
                    except Exception as e:
                        console.print(f"[yellow]Warning: Could not get M3 status: {e}[/yellow]\n")
            
            # Use CLI progress renderer with a gentle heartbeat until completion
            async def _pulse_progress(renderer: CliProgressRenderer):
                current = 0
                while current < 95:
                    progress_obj = ProcessingProgress(
                        stage="transcription",
                        current_step=current,
                        total_steps=100,
                        status="running",
                        message="Transcribing...",
                        start_time=0.0,
                        elapsed_time=0.0,
                    )
                    renderer.on_processing_progress("Transcribing...", progress_obj)
                    await asyncio.sleep(0.3)
                    current += 1
            
            async with asyncio.TaskGroup() as tg:
                with CliProgressRenderer(console) as renderer:
                    renderer.track_processing("Transcribing...", total_steps=100)
                    pulse_task = tg.create_task(_pulse_progress(renderer))
                    
                    if version == 'multi-pass':
                        # Use multi-pass pipeline (run in executor since it's not async)
                        loop = asyncio.get_event_loop()
                        result_dict = await loop.run_in_executor(
                            None,
                            lambda: pipeline.transcribe_with_parallel_processing(
                                audio_path=file_path,
                                speaker_diarization=diarize,
                                domain=domain
                            )
                        )
                        
                        # Convert to expected result format for display
                        from ...services.media_types import TranscriptionResult
                        result = TranscriptionResult(
                            text_content=" ".join([seg.get("text", "") for seg in result_dict.get("segments", [])]),
                            accuracy_estimate=result_dict.get("confidence", 0.9),  # Use average confidence
                            processing_time_ms=result_dict.get("processing_time", 0.0) * 1000,
                            quality_warnings=[],
                            segments=result_dict.get("segments", [])
                        )
                    else:
                        # Use legacy transcription service
                        result = await transcription_service.transcribe_audio(file_path, config)
                    
                    # Complete progress to 100%
                    renderer.on_processing_progress(
                        "Transcribing...",
                        ProcessingProgress(
                            stage="transcription",
                            current_step=100,
                            total_steps=100,
                            status="complete",
                            message="Completed",
                            start_time=0.0,
                            elapsed_time=0.0,
                        ),
                    )
                    
                    # Cancel pulse if still running
                    try:
                        pulse_task.cancel()
                    except Exception:
                        pass
                
                # Check accuracy threshold
                if result.accuracy_estimate and result.accuracy_estimate < min_accuracy / 100.0:
                    console.print(f"[yellow]Warning: Accuracy {result.accuracy_estimate * 100:.1f}% below threshold {min_accuracy}%[/yellow]")
                
                # Display results
                if result.text_content:
                    if output_format == 'json':
                        output_data = {
                            'text': result.text_content,
                            'accuracy': result.accuracy_estimate * 100,
                            'processing_time': result.processing_time_ms / 1000.0,
                            'word_count': len(result.text_content.split()),
                            'quality_warnings': result.quality_warnings
                        }
                        
                        # Add multi-pass specific info
                        if version == 'multi-pass' and hasattr(result, 'segments'):
                            output_data['segments'] = result.segments
                            output_data['pipeline'] = 'multi-pass'
                            if domain:
                                output_data['domain'] = domain
                            output_data['diarization_enabled'] = diarize
                            
                        console.print(json.dumps(output_data, indent=2))
                    else:
                        console.print(f"[bold green]✓[/bold green] Transcription completed!")
                        if version == 'multi-pass':
                            console.print(f"[dim]Pipeline: Multi-pass transcription[/dim]")
                        else:
                            # Show M3 optimization info for v1/v2
                            if hasattr(result, 'metadata') and result.metadata and 'm3_optimizations' in result.metadata:
                                m3_info = result.metadata['m3_optimizations']
                                console.print(f"[dim]Pipeline: v{version} with M3 optimizations[/dim]")
                                console.print(f"[dim]Device: {m3_info.get('device_used', 'unknown')}[/dim]")
                                console.print(f"[dim]M3 Preprocessing: {'✓' if m3_info.get('m3_preprocessing_enabled') else '✗'}[/dim]")
                                console.print(f"[dim]Hardware Acceleration: {'✓' if m3_info.get('hardware_acceleration_used') else '✗'}[/dim]")
                        console.print(f"[cyan]Accuracy:[/cyan] {result.accuracy_estimate * 100:.1f}%")
                        console.print(f"[cyan]Processing time:[/cyan] {result.processing_time_ms / 1000.0:.1f}s")
                        
                        # Show segment count for multi-pass
                        if version == 'multi-pass' and hasattr(result, 'segments'):
                            console.print(f"[cyan]Segments:[/cyan] {len(result.segments)}")
                            if diarize:
                                speakers = set(seg.get('speaker', 'UNKNOWN') for seg in result.segments if seg.get('speaker'))
                                console.print(f"[cyan]Speakers detected:[/cyan] {len(speakers)} ({', '.join(sorted(speakers))})")
                        
                        console.print(f"[cyan]Text:[/cyan]")
                        console.print(result.text_content)
                        
                        if result.quality_warnings:
                            console.print(f"\n[yellow]Quality Warnings:[/yellow]")
                            for warning in result.quality_warnings:
                                console.print(f"  • {warning}")
                else:
                    console.print(f"[red]✗[/red] Transcription failed!")
        
        except Exception as e:
            console.print(f"[red]Error: {str(e)}[/red]")
            raise click.Abort()
    
    asyncio.run(_transcribe())


def find_media_files(directory: Path) -> list:
    """Find all media files in a directory recursively."""
    media_extensions = get_media_extensions()
    media_files = []
    
    for file_path in directory.rglob('*'):
        if file_path.is_file() and file_path.suffix.lower() in media_extensions:
            media_files.append(file_path)
    
    return media_files


def validate_media_directory(directory: Path) -> bool:
    """Validate that a directory contains media files."""
    media_files = find_media_files(directory)
    return len(media_files) > 0
