"""Whisper transcription using faster-whisper."""

import os
import tempfile

import numpy as np
from faster_whisper import WhisperModel
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn


console = Console()


class Transcriber:
    """Transcribes audio using faster-whisper."""
    
    def __init__(self, model_size="base", device="auto", compute_type="default"):
        """Initialize transcriber with Whisper model.
        
        Args:
            model_size: Model size (tiny, base, small, medium, large-v1, large-v2, large-v3)
            device: Device to use (cpu, cuda, auto)
            compute_type: Compute type (int8, int8_float16, int16, float16, float32, default)
        """
        self.model_size = model_size
        self.device = device
        self.compute_type = compute_type
        self.model = None
        
    def load_model(self):
        """Load the Whisper model."""
        if self.model is not None:
            return
            
        console.print(f"[dim]Loading Whisper model ({self.model_size}) on {self.device}...[/dim]")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task("Downloading/loading model...", total=None)
            
            try:
                self.model = WhisperModel(
                    self.model_size,
                    device=self.device,
                    compute_type=self.compute_type,
                    download_root=os.path.expanduser("~/.cache/whisper")
                )
                progress.update(task, description="[green]Model loaded successfully![/green]")
            except Exception as e:
                error_msg = str(e)
                # Check if CUDA-related error
                if "cuda" in error_msg.lower() or "cublas" in error_msg.lower() or "libcuda" in error_msg.lower():
                    console.print(f"[yellow]GPU/CUDA error: {e}[/yellow]")
                    console.print("[dim]Falling back to CPU mode...[/dim]")
                    try:
                        self.model = WhisperModel(
                            self.model_size,
                            device="cpu",
                            compute_type="int8",
                            download_root=os.path.expanduser("~/.cache/whisper")
                        )
                        progress.update(task, description="[green]Model loaded successfully on CPU![/green]")
                    except Exception as cpu_e:
                        progress.update(task, description=f"[red]Error loading model on CPU: {cpu_e}[/red]")
                        raise
                else:
                    progress.update(task, description=f"[red]Error loading model: {e}[/red]")
                    raise
                
    def transcribe(self, audio_data):
        """Transcribe audio data to text.
        
        Args:
            audio_data: numpy array of audio samples at 16kHz
            
        Returns:
            str: Transcribed text
        """
        if self.model is None:
            self.load_model()
            
        # Convert audio to the format expected by faster-whisper
        # Ensure audio is float32 and normalized to [-1, 1]
        if audio_data.dtype != np.float32:
            audio_data = audio_data.astype(np.float32)
            
        # Flatten if stereo (take first channel)
        if len(audio_data.shape) > 1:
            audio_data = audio_data[:, 0]
            
        console.print("[dim]Transcribing audio...[/dim]")
        
        try:
            segments, info = self.model.transcribe(
                audio_data,
                language=None,  # Auto-detect language
                task="transcribe"
            )
            
            # Collect all segments into text
            transcript_parts = []
            for segment in segments:
                transcript_parts.append(segment.text)
                
            transcript = " ".join(transcript_parts).strip()
            
            if transcript:
                console.print(f"[green]Detected language: {info.language} (probability: {info.language_probability:.2f})[/green]")
                return transcript
            else:
                console.print("[yellow]No speech detected in audio[/yellow]")
                return ""
                
        except Exception as e:
            console.print(f"[red]Error during transcription: {e}[/red]")
            return ""
