# audio_processor.py - Final definitive fix for Numba caching issue

import os
from dotenv import load_dotenv
import tempfile
from typing import List, Dict, Optional, Tuple
import json
import traceback

# --- DEFINITIVE NUMBA FIX ---
# This MUST be done BEFORE importing librosa
os.environ["NUMBA_CACHE_DIR"] = "/tmp"

# Now, import librosa safely
import librosa
# --- END OF FIX ---

from google import genai
from google.genai import types

class AudioProcessor:
    def __init__(self):
        self.client = None
        self.init_error = None
        self._initialize_gemini()
    
    def _initialize_gemini(self):
        try:
            load_dotenv()
            api_key = os.getenv("GEMINI_API_KEY")
            if not api_key:
                raise ValueError("GEMINI_API_KEY not found in environment variables.")
            self.client = genai.Client(api_key=api_key)
        except Exception as e:
            self.init_error = f"--- FATAL ERROR during Gemini Init: {str(e)} ---"
            self.client = None
    
    def transcribe_audio(self, audio_file_path: str) -> Tuple[Optional[str], Optional[str]]:
        """
        Transcribes audio. Returns (text, error_message).
        """
        if self.init_error:
            return None, self.init_error
        if not self.client:
            return None, "--- ERROR: Gemini client is not available. ---"
        
        try:
            if not os.path.exists(audio_file_path):
                return None, f"--- ERROR: Audio file for transcription not found at: {audio_file_path} ---"
            
            with open(audio_file_path, 'rb') as f:
                audio_bytes = f.read()
            
            file_ext = os.path.splitext(audio_file_path)[1].lower()
            mime_type = {'mp3': 'audio/mpeg', 'wav': 'audio/wav', 'm4a': 'audio/mp4'}.get(file_ext, 'audio/mpeg')

            response = self.client.models.generate_content(
                model="gemini-2.5-flash",
                contents=[types.Part.from_bytes(data=audio_bytes, mime_type=mime_type), "Transcribe this audio."]
            )
            
            if response and response.text:
                return response.text.strip(), None
            else:
                return None, "--- WARNING: Gemini returned an empty response for transcription. ---"
                
        except Exception as e:
            error_msg = f"--- FATAL ERROR during transcription: {traceback.format_exc()} ---"
            return None, error_msg

    def get_audio_duration(self, audio_file_path: str) -> Tuple[Optional[float], Optional[str]]:
        """
        Gets audio duration. Returns (duration, error_message).
        """
        try:
            if not os.path.exists(audio_file_path):
                return None, f"--- ERROR: Audio file for duration not found at: {audio_file_path} ---"
            
            duration = librosa.get_duration(path=audio_file_path)
            if duration is None or duration < 0.1:
                return None, f"--- ERROR: librosa returned an invalid duration: {duration}s ---"
            return duration, None
        except Exception as e:
            error_msg = f"--- FATAL ERROR getting audio duration with librosa: {traceback.format_exc()} ---"
            return None, error_msg
    
    def get_word_timestamps(self, audio_file_path: str) -> Tuple[List[Dict], List[str]]:
        """
        Generates timestamps. Returns (timestamps, log_messages).
        """
        logs = ["--- INFO: Starting get_word_timestamps... ---"]
        
        transcription, error = self.transcribe_audio(audio_file_path)
        if error:
            logs.append(error)
            return [], logs
        logs.append(f"--- DEBUG: Transcription successful. Text: '{transcription[:50]}...'")

        audio_duration, error = self.get_audio_duration(audio_file_path)
        if error:
            logs.append(error)
            return [], logs
        logs.append(f"--- DEBUG: Audio duration successful. Duration: {audio_duration:.2f}s")
        
        words = transcription.split()
        if not words:
            logs.append("--- WARNING: Transcription resulted in zero words. ---")
            return [], logs
        
        logs.append(f"--- INFO: Distributing {len(words)} words across the duration. ---")
        word_timestamps = []
        total_words = len(words)
        usable_duration = max(0, audio_duration - 1.0)
        
        for i, word in enumerate(words):
            start_time = 0.5 + (i * (usable_duration / total_words))
            end_time = 0.5 + ((i + 1) * (usable_duration / total_words))
            word_timestamps.append({'word': word.strip(), 'start': round(start_time, 3), 'end': round(end_time, 3)})
        
        logs.append(f"--- SUCCESS: Generated {len(word_timestamps)} word timestamps. ---")
        return word_timestamps, logs