from smolagents import CodeAgent,  WikipediaSearchTool, LiteLLMModel, tool, load_tool # HfApiModel, OpenAIServerModel
import asyncio
import os
import yaml

# Simulated additional tools (implementation depends on external APIs or setup)
@tool
class GoogleSearchTool:
    """Tool for performing Google searches using Custom Search JSON API"""
    def __init__(self):
        self.api_key = os.environ.get("GOOGLE_API_KEY")
        self.cse_id = os.environ.get("GOOGLE_CSE_ID")
        if not self.api_key or not self.cse_id:
            raise ValueError("GOOGLE_API_KEY and GOOGLE_CSE_ID must be set in environment variables.")

    def search(self, query: str) -> str:
                """Perform a Google search query
        Args:
            query: Search query string
        Returns:
            Simulated search results (replace with actual API call)
        """
        
        return f"Google search results for '{query}' (simulated)."

@tool
class ImageAnalysisTool:
    """Tool for analyzing images using computer vision"""
    def analyze(self, image_path: str) -> str:
        # Placeholder: Use Google Vision API or similar in real implementation
        """Analyze an image file
        Args:
            image_path: Path to image file
        Returns:
            Simulated image analysis results
        """
        return f"Analyzed image at '{image_path}' (simulated description)."

@tool
class LocalFileAudioTool:
    """Tool for transcribing audio files"""
    def transcribe(self, file_path: str) -> str:
        # Placeholder: Use speech recognition library like SpeechRecognition in real setup
        """Transcribe audio from file
        Args:
            file_path: Path to audio file
        Returns:
            Simulated transcription text
        """    
        return f"Transcribed audio from '{file_path}' (simulated transcription)."

class MagAgent:
    def __init__(self):
        """Initialize the MagAgent with search tools."""
        print("Initializing MagAgent with search tools...")
        model = LiteLLMModel(
            model_id="gemini/gemini-2.0-flash-lite",
            api_key= os.environ.get("GEMINI_KEY"),
            max_tokens=8192
        )

        # Load prompt templates
        with open("prompts.yaml", 'r') as stream:
            prompt_templates = yaml.safe_load(stream)
        
        self.agent = CodeAgent(
            model= model,
            tools=[
                GoogleSearchTool(),
                WikipediaSearchTool(),
                ImageAnalysisTool(),
                LocalFileAudioTool()
            ]
        )
        print("MagAgent initialized.")

    async def __call__(self, question: str) -> str:
        """Process a question asynchronously using the MagAgent."""
        print(f"MagAgent received question (first 50 chars): {question[:50]}...")
        try:
            # Define a task with fallback search logic
            task = (
                f"Answer the following question accurately and concisely: {question}\n"
            )
            response = await asyncio.to_thread(
                self.agent.run,
                task=task
            )

            # Ensure response is a string, fixing the integer error
            response = str(response) if response is not None else "No answer found."
       
            if not response or "No Wikipedia page found" in response:
                # Fallback response if search fails
                response = "Unable to retrieve exact data. Please refine the question or check external sources."
            print(f"MagAgent response: {response[:50]}...")
            return response
        except Exception as e:
            error_msg = f"Error processing question: {str(e)}. Check API key or network connectivity."
            print(error_msg)
            return error_msg