"""Tools for the simple agent.

This module provides the internet search tool using Tavily API.
"""

import os
from typing import Literal

from dotenv import load_dotenv
from langchain_core.tools import tool
from tavily import TavilyClient

# Load environment variables
load_dotenv(override=True)

# Initialize Tavily client
tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API_KEY"))


@tool(parse_docstring=True)
def internet_search(
    query: str,
    max_results: int = 5,
    topic: Literal["general", "news", "finance"] = "general",
    include_raw_content: bool = False,
):
    """Run internet search using Tavily API.

    This is a tool function for web search that wraps Tavily's search functionality.

    Args:
        query: Search query string, e.g. "Python async programming tutorial"
        max_results: Maximum number of results to return, default is 5
        topic: Search topic type, options are "general", "news", or "finance"
        include_raw_content: Whether to include raw webpage content, default is False

    Returns:
        Search results dictionary containing titles, URLs, summaries, etc.
    """
    try:
        result = tavily_client.search(
            query,
            max_results=max_results,
            include_raw_content=include_raw_content,
            topic=topic,
        )
        return result
    except Exception as e:
        # Exception handling: return error message instead of raising exception
        # This allows the LLM to understand the error and try other strategies
        return {"error": f"Search failed: {str(e)}"}