import os
import re
import requests
import serpapi
from smolagents import CodeAgent, ToolCallingAgent, WebSearchTool, tool
from smolagents import InferenceClientModel
from dotenv import load_dotenv
from markdownify import markdownify
from requests.exceptions import RequestException
from llama_index.core import VectorStoreIndex, download_loader

search_cache = {}
webpage_cache = {}

def initialize_agent():
    # Load environment variables from .env file
    load_dotenv()

    # 1. Load the model
    # Make sure to set HF_TOKEN in your Hugging Face Space secrets
    model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
    try:
        model = InferenceClientModel(model_id=model_name, token=os.environ.get("HF_TOKEN"), provider="together")
    except Exception as e:
        print(f"Error loading model: {e}")
        model = None

    # 2. Define the tools
    @tool
    def query_webpage(url: str, query: str) -> str:
        """Queries a webpage at the given URL to find specific information and returns a concise answer.

        Args:
            url: The URL of the webpage to query.
            query: The specific question to ask about the content of the webpage.

        Returns:
            A concise answer to the query based on the webpage's content, or an error message.
        """
        if (url, query) in webpage_cache:
            return webpage_cache[(url, query)]
        try:
            BeautifulSoupWebReader = download_loader("BeautifulSoupWebReader")
            loader = BeautifulSoupWebReader()
            documents = loader.load_data(urls=[url])
            index = VectorStoreIndex.from_documents(documents)
            query_engine = index.as_query_engine()
            response = query_engine.query(query)
            webpage_cache[(url, query)] = str(response)
            return str(response)

        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"

    @tool
    def google_search(query: str) -> str:
        """Searches Google for the given query and returns the results.

        Args:
            query: The query to search for.

        Returns:
            The search results, or an error message if the search fails.
        """
        if query in search_cache:
            return search_cache[query]
        try:
            client = serpapi.Client(api_key=os.environ.get("SERPAPI_API_KEY"))
            results = client.search(q=query, engine="google")
            if "ai_overview" in results:
                ai_overview = results["ai_overview"]
                output = ""
                for block in ai_overview.get("text_blocks", []):
                    if block["type"] == "paragraph":
                        output += block["snippet"] + "\n\n"
                    elif block["type"] == "heading":
                        output += f"### {block['snippet']}\n\n"
                    elif block["type"] == "list":
                        for item in block["list"]:
                            output += f"- **{item['title']}** {item['snippet']}\n"
                        output += "\n"
                if "references" in ai_overview:
                    output += "\n**References:**\n"
                    for ref in ai_overview["references"]:
                        output += f"- [{ref['title']}]({ref['link']})\n"
                search_cache[query] = output
                return output
            elif "organic_results" in results:
                result = str(results["organic_results"])
                search_cache[query] = result
                return result
            else:
                return "No results found."
        except Exception as e:
            return f"Error performing Google search: {str(e)}"

    # 3. Define the agents
    if model:
        web_agent = ToolCallingAgent(
            tools=[WebSearchTool(), query_webpage, google_search],
            model=model,
            max_steps=10,
            name="web_search_agent",
            description="Runs web searches for you.",
        )

        manager_agent = CodeAgent(
            tools=[],
            model=model,
            managed_agents=[web_agent],
            additional_authorized_imports=["time", "numpy", "pandas", "requests", "serpapi", "llama_index", "beautifulsoup4", "markdownify", "lxml", "json", "urllib.parse"],
            instructions='''You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the a new line and the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 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.'''
        )
        return manager_agent
    else:
        return None