transcribe_audio.py

"""
Transcribe podcast audio files using Groq Whisper Large V3 Turbo API.
Handles multilingual transcription and saves results with timestamps.
"""

import json
import os
import re
import time
import tempfile
from pathlib import Path
from groq import Groq
import ffmpeg
import config

def load_metadata():
    """Load metadata.json if it exists, create if it doesn't."""
    if config.METADATA_FILE.exists():
        try:
            with open(config.METADATA_FILE, 'r', encoding='utf-8') as f:
                data = json.load(f)
                # Ensure episodes key exists
                if 'episodes' not in data:
                    data['episodes'] = {}
                return data
        except json.JSONDecodeError:
            print("Warning: metadata.json is corrupted. Creating new one.")
            return {"episodes": {}}
        except Exception as e:
            print(f"Error loading metadata: {e}")
    
    return {"episodes": {}}

def save_metadata(metadata):
    """Save metadata to JSON file."""
    with open(config.METADATA_FILE, 'w', encoding='utf-8') as f:
        json.dump(metadata, f, indent=2, ensure_ascii=False)

def get_file_size_mb(file_path):
    """Get file size in MB."""
    return os.path.getsize(file_path) / (1024 * 1024)

def get_audio_duration(file_path):
    """Get audio duration in seconds using ffprobe."""
    try:
        probe = ffmpeg.probe(file_path)
        duration = float(probe['streams'][0].get('duration', 0))
        if duration == 0:
            # Try format duration
            duration = float(probe['format'].get('duration', 0))
        return duration
    except Exception as e:
        print(f"  Warning: Could not get audio duration: {e}")
        return 0

def split_audio_file(audio_path, max_size_mb=config.GROQ_MAX_FILE_SIZE_MB, overlap_seconds=config.AUDIO_CHUNK_OVERLAP_SECONDS):
    """Split large audio file into smaller chunks with aggressive compression."""
    file_size_mb = get_file_size_mb(audio_path)
    
    if file_size_mb <= max_size_mb:
        return None, []  # No need to split
    
    duration = get_audio_duration(audio_path)
    if duration == 0:
        print(f"  Warning: Could not determine duration, cannot split audio")
        return None, []
    
    # Use more conservative chunk size (50% of max to account for encoding overhead)
    target_chunk_size_mb = max_size_mb * 0.5
    estimated_chunk_duration = (target_chunk_size_mb / file_size_mb) * duration
    
    # Calculate number of chunks needed (add more chunks for safety)
    num_chunks = max(int(duration / estimated_chunk_duration) + 2, int(file_size_mb / target_chunk_size_mb) + 2)
    chunk_duration = duration / num_chunks
    
    print(f"  Splitting {file_size_mb:.1f}MB file ({duration:.0f}s) into ~{num_chunks} chunks (~{chunk_duration:.0f}s each)")
    
    # Create temp directory for chunks
    temp_dir = tempfile.mkdtemp(prefix='audio_chunks_')
    chunk_files = []
    
    try:
        for i in range(num_chunks):
            start_time = i * chunk_duration - (overlap_seconds if i > 0 else 0)
            start_time = max(0, start_time)
            end_time = min((i + 1) * chunk_duration + (overlap_seconds if i < num_chunks - 1 else 0), duration)
            actual_duration = end_time - start_time
            
            chunk_path = os.path.join(temp_dir, f'chunk_{i:03d}.mp3')
            
            # Extract chunk using ffmpeg with aggressive compression
            try:
                stream = ffmpeg.input(audio_path, ss=start_time, t=actual_duration)
                # Use more aggressive compression: 64k bitrate, mono, lower sample rate
                stream = ffmpeg.output(
                    stream, 
                    chunk_path, 
                    acodec='libmp3lame',
                    audio_bitrate='64k',  # Lower bitrate
                    ac=1,  # Mono
                    ar=22050,  # Lower sample rate
                    **{'q:a': '5'}  # Higher compression
                )
                ffmpeg.run(stream, overwrite_output=True, quiet=True, capture_stderr=True)
                
                if not os.path.exists(chunk_path):
                    print(f"  Warning: Chunk {i} file was not created")
                    continue
                
                chunk_size_mb = get_file_size_mb(chunk_path)
                
                # If chunk is still too large, try even more aggressive compression or split further
                if chunk_size_mb > max_size_mb:
                    print(f"  Warning: Chunk {i} is {chunk_size_mb:.1f}MB (too large), trying more aggressive compression...")
                    
                    # Try ultra-compressed version
                    compressed_path = os.path.join(temp_dir, f'chunk_{i:03d}_compressed.mp3')
                    try:
                        stream = ffmpeg.input(audio_path, ss=start_time, t=actual_duration)
                        stream = ffmpeg.output(
                            stream,
                            compressed_path,
                            acodec='libmp3lame',
                            audio_bitrate='32k',  # Very low bitrate
                            ac=1,
                            ar=16000  # Very low sample rate
                        )
                        ffmpeg.run(stream, overwrite_output=True, quiet=True, capture_stderr=True)
                        
                        if os.path.exists(compressed_path):
                            compressed_size_mb = get_file_size_mb(compressed_path)
                            if compressed_size_mb <= max_size_mb:
                                # Remove original, use compressed
                                if os.path.exists(chunk_path):
                                    os.remove(chunk_path)
                                os.rename(compressed_path, chunk_path)
                                chunk_size_mb = compressed_size_mb
                                print(f"  ✓ Chunk {i} compressed to {chunk_size_mb:.1f}MB")
                            else:
                                print(f"  ✗ Chunk {i} still too large ({compressed_size_mb:.1f}MB) after compression, skipping")
                                if os.path.exists(compressed_path):
                                    os.remove(compressed_path)
                                continue
                    except Exception as e:
                        print(f"  Warning: Failed to recompress chunk {i}: {e}")
                        if os.path.exists(compressed_path):
                            os.remove(compressed_path)
                        continue
                
                # Chunk is good, add to list
                chunk_files.append({
                    'path': chunk_path,
                    'start_time': start_time,
                    'end_time': end_time,
                    'index': i,
                    'size_mb': chunk_size_mb
                })
                print(f"  ✓ Chunk {i}: {chunk_size_mb:.1f}MB ({actual_duration:.0f}s)")
                
            except Exception as e:
                print(f"  Warning: Failed to create chunk {i}: {e}")
                if os.path.exists(chunk_path):
                    os.remove(chunk_path)
                continue
        
        if not chunk_files:
            print(f"  Error: No chunks were successfully created")
            return None, []
        
        print(f"  Successfully created {len(chunk_files)} chunks")
        return temp_dir, chunk_files
    
    except Exception as e:
        print(f"  Error splitting audio: {e}")
        import traceback
        traceback.print_exc()
        # Cleanup on error
        import shutil
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        return None, []

def transcribe_audio_chunk(chunk_path, chunk_start_time, client):
    """Transcribe a single audio chunk."""
    try:
        # Determine MIME type
        ext = Path(chunk_path).suffix.lower()
        mime_types = {
            '.mp3': 'audio/mpeg',
            '.m4a': 'audio/mp4',
            '.wav': 'audio/wav',
            '.mp4': 'audio/mp4'
        }
        mime_type = mime_types.get(ext, 'audio/mpeg')
        
        with open(chunk_path, 'rb') as audio_file:
            transcription = client.audio.transcriptions.create(
                file=(os.path.basename(chunk_path), audio_file, mime_type),
                model=config.GROQ_MODEL,
                response_format="verbose_json",
                language=None
            )
        
        return transcription
        
    except Exception as e:
        print(f"  Error transcribing chunk: {e}")
        return None

def merge_chunk_transcriptions(chunk_transcriptions, chunk_files):
    """Merge transcriptions from multiple chunks with proper timestamps."""
    merged = {
        'text': '',
        'language': None,
        'segments': []
    }
    
    for chunk_data, chunk_file_info in zip(chunk_transcriptions, chunk_files):
        if chunk_data is None:
            continue
        
        chunk_start_time = chunk_file_info['start_time']
        
        # Extract data from chunk transcription
        if isinstance(chunk_data, dict):
            text = chunk_data.get('text', '')
            language = chunk_data.get('language', None)
            segments_data = chunk_data.get('segments', [])
        else:
            text = getattr(chunk_data, 'text', '')
            language = getattr(chunk_data, 'language', None)
            segments_data = getattr(chunk_data, 'segments', []) or []
        
        # Set language from first chunk
        if merged['language'] is None and language:
            merged['language'] = language
        
        # Merge text
        if text:
            merged['text'] += text + ' '
        
        # Merge segments with adjusted timestamps
        for segment in segments_data:
            if isinstance(segment, dict):
                seg_start = segment.get('start', 0.0)
                seg_end = segment.get('end', 0.0)
                seg_text = segment.get('text', '')
            else:
                seg_start = getattr(segment, 'start', 0.0)
                seg_end = getattr(segment, 'end', 0.0)
                seg_text = getattr(segment, 'text', '')
            
            # Adjust timestamps to absolute time
            merged['segments'].append({
                'start': chunk_start_time + seg_start,
                'end': chunk_start_time + seg_end,
                'text': seg_text
            })
    
    merged['text'] = merged['text'].strip()
    
    # Sort segments by start time
    merged['segments'].sort(key=lambda x: x['start'])
    
    return merged

def transcribe_audio_file(audio_path, client):
    """Transcribe a single audio file using Groq Whisper API. Handles large files by chunking."""
    try:
        file_size_mb = get_file_size_mb(audio_path)
        print(f"  File size: {file_size_mb:.1f}MB")
        
        # Check if file needs to be split
        if file_size_mb > config.GROQ_MAX_FILE_SIZE_MB:
            print(f"  File exceeds {config.GROQ_MAX_FILE_SIZE_MB}MB limit, splitting...")
            
            temp_dir, chunk_files = split_audio_file(audio_path)
            
            if not chunk_files:
                print(f"  Error: Failed to split audio file")
                return None
            
            print(f"  Transcribing {len(chunk_files)} chunks...")
            
            chunk_transcriptions = []
            try:
                for i, chunk_info in enumerate(chunk_files, 1):
                    print(f"    Chunk {i}/{len(chunk_files)} (starts at {chunk_info['start_time']:.0f}s)...")
                    chunk_trans = transcribe_audio_chunk(
                        chunk_info['path'],
                        chunk_info['start_time'],
                        client
                    )
                    chunk_transcriptions.append(chunk_trans)
                    time.sleep(0.5)  # Rate limiting
                
                # Merge transcriptions
                merged = merge_chunk_transcriptions(chunk_transcriptions, chunk_files)
                
                # Clean up temp files
                import shutil
                if temp_dir and os.path.exists(temp_dir):
                    shutil.rmtree(temp_dir)
                
                return merged
                
            except Exception as e:
                print(f"  Error during chunked transcription: {e}")
                # Clean up on error
                import shutil
                if temp_dir and os.path.exists(temp_dir):
                    shutil.rmtree(temp_dir)
                return None
        
        else:
            # File is small enough, transcribe directly
            print(f"  Uploading {os.path.basename(audio_path)}...")
            
            ext = Path(audio_path).suffix.lower()
            mime_types = {
                '.mp3': 'audio/mpeg',
                '.m4a': 'audio/mp4',
                '.wav': 'audio/wav',
                '.mp4': 'audio/mp4'
            }
            mime_type = mime_types.get(ext, 'audio/mpeg')
            
            with open(audio_path, 'rb') as audio_file:
                transcription = client.audio.transcriptions.create(
                    file=(os.path.basename(audio_path), audio_file, mime_type),
                    model=config.GROQ_MODEL,
                    response_format="verbose_json",
                    language=None
                )
            
            return transcription
        
    except Exception as e:
        print(f"  Error transcribing {audio_path}: {e}")
        import traceback
        traceback.print_exc()
        return None

def save_transcript(episode_id, transcription_data):
    """Save transcript to JSON file."""
    transcript_path = config.TRANSCRIPTS_DIR / f"{episode_id}.json"
    
    # Handle different response formats
    if isinstance(transcription_data, dict):
        # If it's already a dict (from JSON response)
        text = transcription_data.get('text', '')
        language = transcription_data.get('language', None)
        segments_data = transcription_data.get('segments', [])
        words_data = transcription_data.get('words', [])
    else:
        # If it's an object (from SDK response)
        text = getattr(transcription_data, 'text', '')
        language = getattr(transcription_data, 'language', None)
        segments_data = getattr(transcription_data, 'segments', []) or []
        words_data = getattr(transcription_data, 'words', []) or []
    
    # Format transcript data
    transcript = {
        'text': text or '',
        'language': language,
        'segments': []
    }
    
    # Extract segments if available
    if segments_data:
        for segment in segments_data:
            if isinstance(segment, dict):
                transcript['segments'].append({
                    'start': segment.get('start', 0.0),
                    'end': segment.get('end', 0.0),
                    'text': segment.get('text', '')
                })
            else:
                transcript['segments'].append({
                    'start': getattr(segment, 'start', 0.0),
                    'end': getattr(segment, 'end', 0.0),
                    'text': getattr(segment, 'text', '')
                })
    elif words_data:
        # If we have word-level timestamps, create segments
        current_segment = {'start': None, 'end': None, 'text': ''}
        for word in words_data:
            if isinstance(word, dict):
                word_start = word.get('start', 0.0)
                word_end = word.get('end', 0.0)
                word_text = word.get('word', '')
            else:
                word_start = getattr(word, 'start', 0.0)
                word_end = getattr(word, 'end', 0.0)
                word_text = getattr(word, 'word', '')
            
            if current_segment['start'] is None:
                current_segment['start'] = word_start
            
            current_segment['end'] = word_end
            current_segment['text'] += word_text + ' '
            
            # Create new segment every ~5 seconds
            if word_end - current_segment['start'] >= 5.0:
                transcript['segments'].append({
                    'start': current_segment['start'],
                    'end': current_segment['end'],
                    'text': current_segment['text'].strip()
                })
                current_segment = {'start': None, 'end': None, 'text': ''}
        
        # Add final segment
        if current_segment['start'] is not None:
            transcript['segments'].append({
                'start': current_segment['start'],
                'end': current_segment['end'],
                'text': current_segment['text'].strip()
            })
    
    # If no segments available, create segments from text
    if not transcript['segments'] and transcript['text']:
        # Split text into sentence-like segments
        sentences = re.split(r'[.!?]+\s+', transcript['text'])
        current_time = 0.0
        for sentence in sentences:
            if sentence.strip():
                # Estimate ~3 words per second for timing
                words = len(sentence.split())
                duration = words / 3.0
                transcript['segments'].append({
                    'start': current_time,
                    'end': current_time + duration,
                    'text': sentence.strip()
                })
                current_time += duration
    
    # Save transcript
    with open(transcript_path, 'w', encoding='utf-8') as f:
        json.dump(transcript, f, indent=2, ensure_ascii=False)
    
    return transcript_path

def transcribe_episodes():
    """Main function to transcribe all non-transcribed episodes."""
    if not config.GROQ_API_KEY:
        print("Error: GROQ_API_KEY not set in environment or config.py")
        print("Set it with: export GROQ_API_KEY='your-key-here'")
        return
    
    # Initialize Groq client (uses OpenAI-compatible API)
    client = Groq(api_key=config.GROQ_API_KEY)
    metadata = load_metadata()
    
    episodes_to_transcribe = []
    
    # Find episodes that need transcription
    for episode_key, episode_data in metadata.get('episodes', {}).items():
        if (episode_data.get('downloaded', False) and 
            not episode_data.get('transcribed', False)):
            
            audio_path = config.BASE_DIR / episode_data.get('audio_path', '')
            if audio_path.exists():
                episodes_to_transcribe.append((episode_key, episode_data, audio_path))
    
    if not episodes_to_transcribe:
        print("No episodes to transcribe.")
        return
    
    print(f"\n=== Transcribing {len(episodes_to_transcribe)} episodes ===\n")
    
    for i, (episode_key, episode_data, audio_path) in enumerate(episodes_to_transcribe, 1):
        print(f"[{i}/{len(episodes_to_transcribe)}] {episode_data.get('title', 'Unknown')}")
        
        # Extract episode ID for filename
        episode_id = episode_key.split('_', 1)[1] if '_' in episode_key else episode_key
        
        # Transcribe audio
        transcription_data = transcribe_audio_file(str(audio_path), client)
        
        if transcription_data:
            # Save transcript
            transcript_path = save_transcript(episode_id, transcription_data)
            
            # Update metadata
            episode_data['transcribed'] = True
            episode_data['transcript_path'] = str(transcript_path.relative_to(config.BASE_DIR))
            
            # Extract language from transcription data
            if isinstance(transcription_data, dict):
                language = transcription_data.get('language')
            else:
                language = getattr(transcription_data, 'language', None)
            if language:
                episode_data['language'] = language
            
            metadata['episodes'][episode_key] = episode_data
            save_metadata(metadata)
            
            print(f"  ✓ Transcript saved: {transcript_path}")
            
            # Rate limiting - Groq allows high throughput but be respectful
            time.sleep(0.5)
        else:
            print(f"  ✗ Failed to transcribe")
    
    print("\n=== Transcription complete ===")

if __name__ == "__main__":
    transcribe_episodes()