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

@tool
def reverse_sentence_tool(reverse_sentence: str) -> str:
    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:
    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:
    if answer and answer[-1] == '.':
        answer = answer[:-1]
    if "St." in answer:
        answer = answer.replace("St.", "Saint")
    return answer

class BasicAgent:
    def __init__(self):
        api_key = os.environ.get("OPENAI_API_KEY", "")
        # Użyj modelu openai, np. gpt-4o lub gpt-3.5-turbo
        self.model = LiteLLMModel(model_id="gpt-4o", api_key=api_key)
        self.agent = CodeAgent(
            tools=[
                DuckDuckGoSearchTool(),
                PythonInterpreterTool(),
                VisitWebpageTool(),
                reverse_sentence_tool,
                get_youtube_transcript,
                check_answer
            ],
            model=self.model
        )
        print("BasicAgent initialized.")

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