import os
import re
import requests
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

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 = "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B"
    try:
        model = InferenceClientModel(model_id=model_name, token=os.environ.get("HF_TOKEN"))
    except Exception as e:
        print(f"Error loading model: {e}")
        model = None

    # 2. Define the tools
    @tool
    def visit_webpage(url: str) -> str:
        """Visits a webpage at the given URL and returns its content as a markdown string.

        Args:
            url: The URL of the webpage to visit.

        Returns:
            The content of the webpage converted to Markdown, or an error message if the request fails.
        """
        try:
            # Send a GET request to the URL
            response = requests.get(url)
            response.raise_for_status()  # Raise an exception for bad status codes

            # Convert the HTML content to Markdown
            markdown_content = markdownify(response.text).strip()

            # Remove multiple line breaks
            markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)

            return markdown_content

        except RequestException as e:
            return f"Error fetching the webpage: {str(e)}"
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"

    # 3. Define the agents
    if model:
        web_agent = ToolCallingAgent(
            tools=[WebSearchTool(), visit_webpage],
            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", "wikipedia"],
            instructions="""You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with 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