import os
from typing import Optional, Dict, List, Any
import requests
import tempfile
import json
import torch
from base_agent import BaseAgent
from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor

class AudioAgent(BaseAgent):
    """
    Agent for transcribing audio/mp3 file.
    """
    def __init__(self, model: str):
        super().__init__(model=model)

    def __call__(self, question: str, file_name: str, **kwargs) -> str:
        print(f"AudioAgent: Transcribing file: {file_name}")
        final_answer = "Unknown"
        
        try:
            if file_name:

                search_query = self.extract_search_query(question)
        
                try:
                    # Set the model ID based on size
                    model_id = "openai/whisper-tiny"
                    
                    # Load model with optimizations for CPU
                    model = AutoModelForSpeechSeq2Seq.from_pretrained(
                        model_id,
                        torch_dtype=torch.float32,  # Use float32 for CPU
                        low_cpu_mem_usage=True,
                        use_safetensors=True
                    )
                    
                    # Load processor
                    processor = AutoProcessor.from_pretrained(model_id)
                    
                    # Create pipeline with small batch size for CPU
                    pipe = pipeline(
                        "automatic-speech-recognition",
                        model=model,
                        tokenizer=processor.tokenizer,
                        feature_extractor=processor.feature_extractor,
                        max_new_tokens=128,
                        chunk_length_s=30,
                        batch_size=1,  # Small batch size for CPU
                        return_timestamps=True,
                        device="cpu",
                    )
                    
                    print("Model loaded successfully.")
        
                    print(f"Starting transcription of {file_name}")
        
                    try:
                        print("Running transcription (this may take some time on CPU)...")
                        
                        # Set task and language options
                        generate_kwargs = {}
                        generate_kwargs["language"] = "english"
                            
                        # Run the pipeline directly on the file
                        result = pipe(
                            file_name,
                            generate_kwargs=generate_kwargs,
                            return_timestamps=True
                        )
                        
                        # Process results based on output format
        
                        transcript = result["text"]
        
                        system_prompt = """
                        You are a research expert tasked with answering factual questions based on a text transcript extracted from an audio file.
                        Analyze the transcript carefully and extract the information needed to answer the question accurately.
                        Provide a concise, accurate answer based only on the provided information.
                        You must provide only the exact answer requested by the user, in the correct type and minimal format, based purely on the provided information.
                        YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
                        If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
                        If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
                        If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
                        If you cannot find the answer based on the provided information, reply only with "Unknown".
                        """
                        
                        user_prompt = f"""
                        Question: {question}
                        
                        Transcript:
                        {transcript}
                        
                        Based on these search results, provide a concise and accurate answer to the question.
                        """
        
                        final_answer = self.llm_call(system_prompt, user_prompt)
        
                        print(f"Final answer: {final_answer}")
                
                        if final_answer.lower() == "unknown":
                            final_answer = self.ask_james(question)                
        
                    except Exception as e:
                        print(f"Transcription failed: {str(e)}")
                        raise
                    
                except Exception as e:
                    print(f"Error loading model: {str(e)}")
                    raise

        except Exception as e:
            print(f"Transcription failed: {str(e)}")
            raise

        return final_answer