from smolagents import CodeAgent, LiteLLMModel, tool, Tool, load_tool, DuckDuckGoSearchTool, WikipediaSearchTool
import asyncio
import os
import re
import pandas as pd
from typing import Optional
from token_bucket import Limiter, MemoryStorage
import yaml
from PIL import Image, ImageOps
import requests
from io import BytesIO
from markdownify import markdownify
import whisper
import time
import shutil
import traceback
from langchain_community.document_loaders import ArxivLoader

@tool
def search_arxiv(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 result.
    
    Args:
        query: The search query.
     Returns:
        str: Formatted search results
    """
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ])
    return {"arxiv_results": formatted_search_docs}

class VisitWebpageTool(Tool):
    name = "visit_webpage"
    description = "Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages."
    inputs = {'url': {'type': 'string', 'description': 'The url of the webpage to visit.'}}
    output_type = "string"

    def forward(self, url: str) -> str:
        try:
            response = requests.get(url, timeout=50)
            response.raise_for_status()
            markdown_content = markdownify(response.text).strip()
            markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
            from smolagents.utils import truncate_content
            return truncate_content(markdown_content, 10000)
        except requests.exceptions.Timeout:
            return "The request timed out. Please try again later or check the URL."
        except requests.exceptions.RequestException as e:
            return f"Error fetching the webpage: 
{str(e)}"
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"

    def __init__(self, *args, **kwargs):
        self.is_initialized = False

class SpeechToTextTool(Tool):
    name = "speech_to_text"
    description = (
        "Converts an audio file to text using OpenAI Whisper."
    )
    inputs = {
        "audio_path": {"type": "string", "description": "Path to audio file (.mp3, .wav)"},
    }
    output_type = "string"

    def __init__(self):
        super().__init__()
        self.model = whisper.load_model("base")

    def forward(self, audio_path: str) -> str:
        if not os.path.exists(audio_path):
            return f"Error: File not found at {audio_path}"
        result = self.model.transcribe(audio_path)
        return result.get("text", "")

class ExcelReaderTool(Tool):
    name = "excel_reader"
    description = """
    This tool reads and processes Excel files (.xlsx, .xls).
    It can extract data, calculate statistics, and perform data analysis on spreadsheets.
    """
    inputs = {
        "excel_path": {
            "type": "string",
            "description": "The path to the Excel file to read",
        },
        "sheet_name": {
            "type": "string",
            "description": "The name of the sheet to read (optional, defaults to first sheet)",
            "nullable": True
        }
    }
    output_type = "string"
    
    def forward(self, excel_path: str, sheet_name: str = None) -> str:
        try:
            if not os.path.exists(excel_path):
                return f"Error: Excel file not found at {excel_path}"
            import pandas as pd
            if sheet_name:
                df = pd.read_excel(excel_path, sheet_name=sheet_name)
            else:
                df = pd.read_excel(excel_path)
            info = {
                "shape": df.shape,
                "columns": list(df.columns),
                "dtypes": df.dtypes.to_dict(),
                "head": df.head(5).to_dict()
            }
            result = f"Excel file: {excel_path}\n"
            result += f"Shape: {info['shape'][0]} rows × {info['shape'][1]} columns\n\n"
            result += "Columns:\n"
            for col in info['columns']:
                result += f"- {col} ({info['dtypes'].get(col)})\n"
            result += "\nPreview (first 5 rows):\n"
            result += df.head(5).to_string()
            return result
        except Exception as e:
            return f"Error reading Excel file: {str(e)}"

class PythonCodeReaderTool(Tool):
    name = "read_python_code"
    description = "Reads a Python (.py) file and returns its content as a string."
    inputs = {
        "file_path": {"type": "string", "description": "The path to the Python file to read"}
    }
    output_type = "string"

    def forward(self, file_path: str) -> str:
        try:
            if not os.path.exists(file_path):
                return f"Error: Python file not found at {file_path}"
            with open(file_path, "r", encoding="utf-8") as file:
                content = file.read()
            return content
        except Exception as e:
            return f"Error reading Python file: {str(e)}"

from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

class RetryDuckDuckGoSearchTool(DuckDuckGoSearchTool):
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type(Exception)
    )
    def forward(self, query: str) -> str:
        return super().forward(query)

class MagAgent:
    def __init__(self, rate_limiter: Optional[Limiter] = None):
        """Initialize the MagAgent with search tools."""
        self.rate_limiter = rate_limiter

        print("Initializing MagAgent with search tools...")
        model = LiteLLMModel(
            model_id="gemini/gemini-2.0-flash",
            api_key=os.environ.get("GEMINI_KEY"),
            max_tokens=8192
        )

        self.agent = CodeAgent(
            model=model,
            tools=[
                RetryDuckDuckGoSearchTool(),
                WikipediaSearchTool(),
                SpeechToTextTool(),
                ExcelReaderTool(),
                VisitWebpageTool(),
                PythonCodeReaderTool(),
                search_arxiv,
            ],
            verbosity_level=2,
            add_base_tools=False,
            max_steps=20
        )
        print("MagAgent initialized.")

    async def __call__(self, question: str, file_path: Optional[str] = None) -> str:
        """Process a question asynchronously using the MagAgent."""
        print(f"MagAgent received question (first 50 chars): {question[:50]}... File path: {file_path}")
        try:
            if self.rate_limiter:
                while not self.rate_limiter.consume(1):
                    print(f"Rate limit reached. Waiting...")
                    await asyncio.sleep(4)
            task = (
                "You are an advanced AI assistant tasked with answering questions from the GAIA benchmark accurately and concisely. Follow these guidelines:\n\n"
                "   - If the question includes direct speech or quoted text (e.g., \"Isn't that hot?\"), treat it as a precise query and preserve the quoted structure in your response.\n\n"
            )
            if file_path:
                task += f"   - The question references an attachment. The attachment file is located at: {file_path}. Use the appropriate tool based on the file extension to process it.\n"
            task += (
                "You are an advanced AI assistant tasked with answering questions from the GAIA benchmark accurately and concisely. Follow these guidelines:\n\n"
                "   - If the question includes direct speech or quoted text (e.g., \"Isn't that hot?\"), treat it as a precise query and preserve the quoted structure in your response.\n\n"
                "   - When processing external data (e.g., YouTube transcripts, web searches), expect potential issues like missing punctuation, inconsistent formatting, or conversational text.\n"
                "   - If the input is ambiguous, prioritize extracting key information relevant to the question.\n\n"
                "   - Provide answers that are concise, accurate, and properly punctuated according to standard English grammar.\n"
                "   - When the answer is a direct quote or direct speech, include the quotation marks in the final answer submitted via `final_answer`. For example, if the answer is \"Extremely.\", submit it as \"\\\"Extremely.\\\"\"\n\n"
                "   - If asked about name of place or city, use full complete name without abbreviations (e.g. use Saint Petersburg instead of St.Petersburg). \n"
                "   - When giving the final answer, output only the direct required result without any extra text like \"Final Answer:\" or explanations.\n"
                f"Answer the following question accurately and concisely: \n {question} \n"
            )
            print(f"Calling agent.run...")
            response = await asyncio.to_thread(self.agent.run, task=task)
            print(f"Agent.run completed.")
            response = str(response)
            if not response:
                print(f"No answer found.")
                response = "No answer found."
            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