"""
Groq Speech-to-Text Transcription Utility

Fast, cost-effective audio transcription using Groq's Whisper API.
Supports both transcription (any language to text) and translation (any language to English).

Recommended Model: whisper-large-v3-turbo (best price/performance with multilingual support)
"""
from __future__ import annotations

import base64
import io
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal

from groq import Groq
from matrx_utils import vcprint


@dataclass
class TranscriptionUsage:
    """Track usage for transcription requests"""
    
    duration_seconds: float
    """Duration of the audio file in seconds"""
    
    model: str
    """Whisper model used for transcription"""
    
    language: str | None = None
    """Language of the audio (ISO-639-1 format like 'en', 'es', 'tr')"""
    
    operation: Literal["transcription", "translation"] = "transcription"
    """Type of operation performed"""
    
    billed_duration: float = 0.0
    """Actual billed duration (minimum 10 seconds)"""
    
    file_size_mb: float = 0.0
    """Size of the audio file in MB"""
    
    response_format: str = "json"
    """Format of the response"""
    
    metadata: dict[str, Any] = field(default_factory=dict)
    """Additional metadata from the transcription"""
    
    def __post_init__(self):
        """Calculate billed duration (minimum 10 seconds per Groq pricing)"""
        self.billed_duration = max(self.duration_seconds, 10.0)
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for storage"""
        return {
            "duration_seconds": self.duration_seconds,
            "billed_duration": self.billed_duration,
            "model": self.model,
            "language": self.language,
            "operation": self.operation,
            "file_size_mb": self.file_size_mb,
            "response_format": self.response_format,
            "metadata": self.metadata,
        }

@dataclass
class TranscriptionResult:
    """Result from a transcription or translation request"""
    
    text: str
    """The transcribed/translated text"""
    
    usage: TranscriptionUsage
    """Usage tracking information"""
    
    segments: list | None = None
    """Detailed segment information (if verbose_json was requested)"""
    
    language: str | None = None
    """Detected language (for transcriptions)"""
    
    duration: float | None = None
    """Duration from response metadata"""
    
    quality_metrics: dict[str, Any] = field(default_factory=dict)
    """Quality metrics (avg_logprob, no_speech_prob, compression_ratio)"""
    
    raw_response: dict[str, Any] | None = None
    """Full raw response from Groq API"""
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for storage"""
        return {
            "text": self.text,
            "usage": self.usage.to_dict(),
            "segments": self.segments,
            "language": self.language,
            "duration": self.duration,
            "quality_metrics": self.quality_metrics,
        }

class GroqTranscription:
    """
    Groq Speech-to-Text Transcription Service
    
    Provides fast, cost-effective audio transcription using Groq's Whisper API.
    
    Features:
    - Transcription: Convert audio in any language to text in that language
    - Translation: Convert audio in any language to English text
    - Multiple Whisper models available
    - Automatic audio preprocessing
    - Usage tracking and quality metrics
    
    Recommended Model: whisper-large-v3-turbo
    """
    
    # Supported file formats
    SUPPORTED_FORMATS = ["flac", "mp3", "mp4", "mpeg", "mpga", "m4a", "ogg", "wav", "webm"]
    
    # Available Whisper models
    MODELS = {
        "whisper-large-v3": "Most accurate, multilingual, error-sensitive applications",
        "whisper-large-v3-turbo": "Best price/performance, multilingual (RECOMMENDED)",
    }
    
    # File size limits (in MB)
    MAX_FILE_SIZE_FREE = 25  # 25 MB for free tier
    MAX_FILE_SIZE_DEV = 100  # 100 MB for dev tier
    
    def __init__(
        self,
        api_key: str | None = None,
        default_model: str = "whisper-large-v3-turbo",
        debug: bool = False,
    ):
        """
        Initialize Groq Transcription client
        
        Args:
            api_key: Groq API key (defaults to GROQ_API_KEY env var)
            default_model: Default Whisper model to use
            debug: Enable debug logging
        """
        self.client = Groq(api_key=api_key or os.environ.get("GROQ_API_KEY"))
        self.default_model = default_model
        self.debug = debug
        
        if self.debug:
            vcprint(
                f"Initialized GroqTranscription with model: {self.default_model}",
                "Groq Transcription",
                color="blue"
            )
    
    def transcribe(
        self,
        audio_source: str | bytes | io.BytesIO,
        model: str | None = None,
        language: str | None = None,
        prompt: str | None = None,
        response_format: Literal["json", "verbose_json", "text"] = "verbose_json",
        temperature: float = 0.0,
        timestamp_granularities: list | None = None,
    ) -> TranscriptionResult:
        """
        Transcribe audio to text in the same language
        
        Args:
            audio_source: File path, URL, base64 data, or bytes
            model: Whisper model to use (defaults to default_model)
            language: Language of the input audio (ISO-639-1 format, e.g., 'en', 'es')
            prompt: Optional prompt to guide transcription style (max 224 tokens)
            response_format: Output format (json, verbose_json, or text)
            temperature: Sampling temperature (0-1, default 0 for deterministic)
            timestamp_granularities: List of ["word", "segment"] for timestamps
        
        Returns:
            TranscriptionResult with text and metadata
        """
        model = model or self.default_model
        
        if self.debug:
            vcprint(
                f"Starting transcription with model: {model}",
                "Groq Transcription",
                color="cyan"
            )
        
        # Prepare audio file
        file_tuple, file_size_mb = self._prepare_audio_file(audio_source)
        
        # Build request parameters
        request_params = {
            "file": file_tuple,
            "model": model,
            "response_format": response_format,
            "temperature": temperature,
        }
        
        if language:
            request_params["language"] = language
        
        if prompt:
            request_params["prompt"] = prompt
        
        if timestamp_granularities and response_format == "verbose_json":
            request_params["timestamp_granularities"] = timestamp_granularities
        
        # Make API call
        try:
            response = self.client.audio.transcriptions.create(**request_params)
            
            if self.debug:
                vcprint(response, "Groq Transcription Response", color="green")
            
            # Parse response
            return self._parse_response(
                response,
                model=model,
                operation="transcription",
                file_size_mb=file_size_mb,
                response_format=response_format,
            )
        
        except Exception as e:
            vcprint(f"Transcription error: {str(e)}", "Groq Transcription Error", color="red")
            raise
    
    def translate(
        self,
        audio_source: str | bytes | io.BytesIO,
        model: str | None = None,
        prompt: str | None = None,
        response_format: Literal["json", "text"] = "json",
        temperature: float = 0.0,
    ) -> TranscriptionResult:
        """
        Translate audio to English text
        
        Args:
            audio_source: File path, URL, base64 data, or bytes
            model: Whisper model to use (defaults to default_model)
            prompt: Optional prompt to guide translation style (max 224 tokens)
            response_format: Output format (json or text)
            temperature: Sampling temperature (0-1, default 0 for deterministic)
        
        Returns:
            TranscriptionResult with English text and metadata
        """
        model = model or self.default_model
        
        if self.debug:
            vcprint(
                f"Starting translation with model: {model}",
                "Groq Translation",
                color="cyan"
            )
        
        # Prepare audio file
        file_tuple, file_size_mb = self._prepare_audio_file(audio_source)
        
        # Build request parameters
        request_params = {
            "file": file_tuple,
            "model": model,
            "response_format": response_format,
            "temperature": temperature,
            "language": "en",  # Translation endpoint only supports English
        }
        
        if prompt:
            request_params["prompt"] = prompt
        
        # Make API call
        try:
            response = self.client.audio.translations.create(**request_params)
            
            if self.debug:
                vcprint(response, "Groq Translation Response", color="green")
            
            # Parse response
            return self._parse_response(
                response,
                model=model,
                operation="translation",
                file_size_mb=file_size_mb,
                response_format=response_format,
            )
        
        except Exception as e:
            vcprint(f"Translation error: {str(e)}", "Groq Translation Error", color="red")
            raise
    
    def _prepare_audio_file(
        self, audio_source: str | bytes | io.BytesIO
    ) -> tuple:
        """
        Prepare audio file for API request
        
        Returns:
            Tuple of (file_tuple, file_size_mb)
        """
        if isinstance(audio_source, str):
            # Could be file path or URL
            if audio_source.startswith(("http://", "https://")):
                # It's a URL - download it first
                from matrx_ai.media import fetch_media
                audio_data = fetch_media(audio_source, target_format="bytes")
                file_tuple = ("audio.wav", audio_data)
                file_size_mb = len(audio_data) / (1024 * 1024)
            elif audio_source.startswith("data:"):
                # It's a base64 data URI
                audio_data = base64.b64decode(audio_source.split(",")[1])
                file_tuple = ("audio.wav", audio_data)
                file_size_mb = len(audio_data) / (1024 * 1024)
            else:
                # It's a file path
                file_path = Path(audio_source)
                if not file_path.exists():
                    raise FileNotFoundError(f"Audio file not found: {audio_source}")
                
                file_size_mb = file_path.stat().st_size / (1024 * 1024)
                file_tuple = (file_path.name, open(file_path, "rb"))
        
        elif isinstance(audio_source, bytes):
            # Raw bytes
            file_tuple = ("audio.wav", audio_source)
            file_size_mb = len(audio_source) / (1024 * 1024)
        
        elif isinstance(audio_source, io.BytesIO):
            # BytesIO object
            audio_data = audio_source.getvalue()
            file_tuple = ("audio.wav", audio_data)
            file_size_mb = len(audio_data) / (1024 * 1024)
        
        else:
            raise ValueError(f"Unsupported audio_source type: {type(audio_source)}")
        
        # Check file size
        if file_size_mb > self.MAX_FILE_SIZE_DEV:
            raise ValueError(
                f"Audio file too large ({file_size_mb:.2f} MB). "
                f"Maximum size is {self.MAX_FILE_SIZE_DEV} MB. "
                "Consider chunking the audio file."
            )
        
        if self.debug:
            vcprint(f"Prepared audio file: {file_size_mb:.2f} MB", color="blue")
        
        return file_tuple, file_size_mb
    
    def _parse_response(
        self,
        response: Any,
        model: str,
        operation: Literal["transcription", "translation"],
        file_size_mb: float,
        response_format: str,
    ) -> TranscriptionResult:
        """Parse API response into TranscriptionResult"""
        
        # Extract text
        if isinstance(response, str):
            text = response
            segments = None
            duration = None
            language = None
            quality_metrics = {}
            raw_response = None
        else:
            # Response is an object
            text = response.text if hasattr(response, "text") else str(response)
            
            # Extract verbose_json fields if available
            segments = getattr(response, "segments", None)
            duration = getattr(response, "duration", None)
            language = getattr(response, "language", None)
            
            # Calculate quality metrics from segments
            quality_metrics = {}
            if segments:
                quality_metrics = self._calculate_quality_metrics(segments)
            
            # Store raw response
            raw_response = response.model_dump() if hasattr(response, "model_dump") else None
        
        # Create usage tracking
        usage = TranscriptionUsage(
            duration_seconds=duration or 0.0,
            model=model,
            language=language,
            operation=operation,
            file_size_mb=file_size_mb,
            response_format=response_format,
            metadata={
                "has_segments": segments is not None,
                "segment_count": len(segments) if segments else 0,
            }
        )
        
        return TranscriptionResult(
            text=text,
            usage=usage,
            segments=segments,
            language=language,
            duration=duration,
            quality_metrics=quality_metrics,
            raw_response=raw_response,
        )
    
    def _calculate_quality_metrics(self, segments: list) -> dict[str, Any]:
        """
        Calculate aggregate quality metrics from segments
        
        Metrics:
        - avg_logprob: Average log probability (closer to 0 = higher confidence)
        - no_speech_prob: Probability of no speech (closer to 0 = definitely speech)
        - compression_ratio: Speech pattern indicator (1.5-2.5 is normal)
        """
        if not segments:
            return {}
        
        total_logprob = 0.0
        total_no_speech = 0.0
        total_compression = 0.0
        count = 0
        
        low_confidence_segments = []
        high_no_speech_segments = []
        unusual_compression_segments = []
        
        for segment in segments:
            if isinstance(segment, dict):
                avg_logprob = segment.get("avg_logprob", 0)
                no_speech_prob = segment.get("no_speech_prob", 0)
                compression_ratio = segment.get("compression_ratio", 0)
            else:
                avg_logprob = getattr(segment, "avg_logprob", 0)
                no_speech_prob = getattr(segment, "no_speech_prob", 0)
                compression_ratio = getattr(segment, "compression_ratio", 0)
            
            total_logprob += avg_logprob
            total_no_speech += no_speech_prob
            total_compression += compression_ratio
            count += 1
            
            # Flag problematic segments
            if avg_logprob < -0.5:  # Low confidence
                low_confidence_segments.append({
                    "id": segment.get("id") if isinstance(segment, dict) else getattr(segment, "id", None),
                    "start": segment.get("start") if isinstance(segment, dict) else getattr(segment, "start", None),
                    "end": segment.get("end") if isinstance(segment, dict) else getattr(segment, "end", None),
                    "avg_logprob": avg_logprob,
                })
            
            if no_speech_prob > 0.5:  # Likely not speech
                high_no_speech_segments.append({
                    "id": segment.get("id") if isinstance(segment, dict) else getattr(segment, "id", None),
                    "start": segment.get("start") if isinstance(segment, dict) else getattr(segment, "start", None),
                    "end": segment.get("end") if isinstance(segment, dict) else getattr(segment, "end", None),
                    "no_speech_prob": no_speech_prob,
                })
            
            if compression_ratio < 1.0 or compression_ratio > 3.0:  # Unusual
                unusual_compression_segments.append({
                    "id": segment.get("id") if isinstance(segment, dict) else getattr(segment, "id", None),
                    "start": segment.get("start") if isinstance(segment, dict) else getattr(segment, "start", None),
                    "end": segment.get("end") if isinstance(segment, dict) else getattr(segment, "end", None),
                    "compression_ratio": compression_ratio,
                })
        
        metrics = {
            "avg_logprob": total_logprob / count if count > 0 else 0,
            "avg_no_speech_prob": total_no_speech / count if count > 0 else 0,
            "avg_compression_ratio": total_compression / count if count > 0 else 0,
            "segment_count": count,
        }
        
        # Add warnings for problematic segments
        if low_confidence_segments:
            metrics["low_confidence_segments"] = low_confidence_segments
            metrics["warning_low_confidence"] = True
        
        if high_no_speech_segments:
            metrics["high_no_speech_segments"] = high_no_speech_segments
            metrics["warning_high_no_speech"] = True
        
        if unusual_compression_segments:
            metrics["unusual_compression_segments"] = unusual_compression_segments
            metrics["warning_unusual_compression"] = True
        
        return metrics
