from smolagents import LiteLLMModel
from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, tool, PythonInterpreterTool
from youtube_transcript_api import YouTubeTranscriptApi
import os

@tool
def reverse_sentence_tool(reverse_sentence: str) -> str:
    """
    Receives a sentence where both the word order and the characters in each word are reversed.
    Returns the sentence with words and order corrected.
    Args:
        reverse_sentence: A sentence with reversed words and reversed word order.
    Returns:
        A sentence in natural reading order.
    """
    inverted_words = reverse_sentence.split(" ")[::-1]
    correct_words = [word[::-1] for word in inverted_words]
    return " ".join(correct_words)

@tool
def get_youtube_transcript(video_url: str) -> str:
    """
    Fetches the transcript from a YouTube video if available.
    Args:
        video_url: Full URL to the YouTube video.
    Returns:
        Transcript text.
    """
    video_id = video_url.split("v=")[-1]
    transcript = YouTubeTranscriptApi.get_transcript(video_id)
    full_text = " ".join([entry['text'] for entry in transcript])
    return full_text

@tool
def check_answer(answer: str) -> str:
    """
    Reviews the answer to check that it meets the requirements specified by the user and modifies it if necessary.
    Args:
        answer (str): The answer of the Agent.
    Returns:
        str: The final answer.
    """
    if answer and answer[-1] == '.':
        answer = answer[:-1]
    if "St." in answer:
        answer = answer.replace("St.", "Saint")
    return answer

class BasicAgent:
    def __init__(self):
        # Odczytaj klucz OpenAI z ENV
        self.api_key = os.getenv("OPENAI_API_KEY")
        # Ustaw preferowany model, np. GPT-4o lub inny OpenAI
        self.model = LiteLLMModel(model_id="gpt-4o", api_key=self.api_key)

        self.agent = CodeAgent(
            tools=[
                DuckDuckGoSearchTool(),
                PythonInterpreterTool(),
                VisitWebpageTool(),
                reverse_sentence_tool,
                get_youtube_transcript,
                check_answer,
            ],
            model=self.model
        )
        print("BasicAgent initialized (OpenAI).")

    def __call__(self, question: str) -> str:
        print(f"Agent received question: {question[:50]}...")
        answer = self.agent.run(question)
        return answer