import os
from smolagents import Tool
from googleapiclient.discovery import build

class YouTubeSearchTool(Tool):
    """
    A tool to search for YouTube videos using the Google Custom Search Engine (CSE) API,
    by specifically restricting results to the youtube.com domain.
    """
    name = "youtube_search"
    description = "Searches for and returns links and snippets for YouTube videos related to a specific query."

    inputs = {
        # Corrected the typo/broken line to read correctly
        "query": {"type": "string", "description": "The search term to find relevant YouTube videos for."}
    }
    output_type = "string"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        # Retrieve credentials from environment variables
        self.api_key = os.getenv("GOOGLE_API_KEY")
        self.cse_id = os.getenv("GOOGLE_CSE_ID")

        if not self.api_key or not self.cse_id:
            # Note: This check relies on environment variables being set up
            raise ValueError("GOOGLE_API_KEY or GOOGLE_CSE_ID secret not found.")

        # Initialize the Google Custom Search service
        self.service = build(
            "customsearch", "v1", developerKey=self.api_key
        )

    def forward(self, query: str) -> str:
        """
        Executes a Google search query, restricting results to youtube.com.

        Args:
            query: The search term provided by the agent.

        Returns:
            A formatted string of search results (max 3), or an error message.
        """
        print(f"Executing YouTube search for: '{query}'")
        
        # Modify the query to ensure results are restricted to YouTube
        youtube_query = f"{query} site:youtube.com"

        try:
            # Execute the search request for up to 3 results
            res = self.service.cse().list(
                q=youtube_query,
                cx=self.cse_id,
                num=3
            ).execute()

            items = res.get('items', [])

            if not items:
                return "XX record info: No YouTube results found."

            search_results = []
            for i, item in enumerate(items):
                search_results.append(
                    f"RESULT {i+1}: '{item.get('title')}'\n"
                    f"CONTENT: {item.get('snippet')}\n"
                    f"SOURCE: {item.get('link')}"
                )

            return "\n\n---SEPARATOR---\n\n".join(search_results)

        except Exception as e:
            return f"Error during YouTube Search API call: {e}"