from dotenv import load_dotenv
from smolagents import tool, PythonInterpreterTool, DuckDuckGoSearchTool, WikipediaSearchTool, VisitWebpageTool, GoogleSearchTool, ToolCallingAgent, CodeAgent, LiteLLMModel
from youtube_transcript_api import YouTubeTranscriptApi
import os
from typing import Dict

load_dotenv()

system_prompt = f"""
You are a general AI assistant that will receive a question. 
First, provide a step-by-step explanation of your reasoning to answer the question, 
providing 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.
**Rules to follow when generating your answer:**
- 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.
"""

@tool
def add(a: int, b: int) -> int:
    #Google docstring style
    """
    Add two numbers together
    Args:
        a: The first number
        b: The second number
    Returns:
        The sum of the two numbers
    """
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """
    Subtract two numbers
    Args:
        a: The first number
        b: The second number
    Returns:
        The difference between the two numbers
    """
    return a - b

@tool
def multiply(a: int, b: int) -> int:
    """
    Multiply two numbers
    Args:
        a: The first number
        b: The second number
    Returns:
        The product of the two numbers
    """
    return a * b

@tool
def divide(a: int, b: int) -> float:
    """
    Divide two numbers
    Args:
        a: The first number
        b: The second number
    Returns:
        The quotient of the two numbers
    """
    return a / b    

@tool
def modulo(a: int, b: int) -> int:
    """
    Calculate the modulo of two numbers
    Args:
        a: The first number
        b: The second number
    Returns:
        The remainder of the division of the two numbers
    """
    return a % b

@tool
def youtube_transcript(url: str) -> Dict[str, str]:
    """
    Get transcript of YouTube video.
    Args:
        url: YouTube video url in ""
    Returns:
        Transcript of the YouTube video
    """    
    video_id = url.partition("https://www.youtube.com/watch?v=")[2]
    transcript = YouTubeTranscriptApi.get_transcript(video_id)
    transcript_text = " ".join([item["text"] for item in transcript])
    return {"youtube_transcript": transcript_text}

class BasicSmolAgent:
    def __init__(self):
        self.api_key = os.getenv("OPENAI_API_KEY")
        self.model = LiteLLMModel(model_id="openai/o4-mini", api_key=self.api_key)
        self.agent = CodeAgent(
            tools=[
                add, subtract, multiply, divide, modulo,
                youtube_transcript,
                DuckDuckGoSearchTool(),
                WikipediaSearchTool(),
                VisitWebpageTool(),
                PythonInterpreterTool(),
                
            ],
            model=self.model
        )
        
    def __call__(self, question: str) -> str:
        print(f"Question: {question}")
        return self.agent.run(question)