import os
import logging
from typing import Literal, Union
from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
from googleapiclient.discovery import build


# Configure basic logging for tools
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- Arithmetic Tool ---
@tool
def perform_calculation(
    a: Union[int, float],
    b: Union[int, float],
    operation: Literal["add", "subtract", "multiply", "divide", "modulus"],
) -> Union[int, float, str]:
    """Performs a specified arithmetic operation on two numbers.

    Args:
        a: The first number (integer or float).
        b: The second number (integer or float).
        operation: The arithmetic operation to perform.
                   Must be one of "add", "subtract", "multiply", "divide", or "modulus".
    """
    logging.info(f"Tool Call: perform_calculation(a={a}, b={b}, operation='{operation}')")
    try:
        if operation == "add":
            result = a + b
        elif operation == "subtract":
            result = a - b
        elif operation == "multiply":
            result = a * b
        elif operation == "divide":
            if b == 0:
                logging.warning("perform_calculation: Division by zero attempted.")
                return "Error: Cannot divide by zero."  
            result = a / b
        elif operation == "modulus":
            if b == 0:
                logging.warning("perform_calculation: Modulus with zero divisor attempted.")
                return "Error: Cannot perform modulus with zero divisor."
            # Optional: enforce integer for modulus for clarity/LLM expectation
            if not isinstance(a, int) or not isinstance(b, int):
                logging.warning("perform_calculation: Modulus with non-integer operands. Results might be float.")
                # You might return a specific error or proceed depending on desired behavior
            result = a % b
        else:
            logging.error(f"perform_calculation: Invalid operation '{operation}' specified.")
            return "Error: Invalid operation specified."
        logging.info(f"perform_calculation: Result of {operation} on {a}, {b} is {result}")
        return result
    except Exception as e:
        logging.error(f"perform_calculation: An unexpected error occurred: {e}", exc_info=True)
        return f"Error during calculation: {e}"


# --- Search & Document Loader Tools ---

@tool
def search_from_wiki(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 relevant sections.
    
    Args:
        query: The search query for Wikipedia."""
    logging.info(f"Tool Call: search_from_wiki(query='{query}')")
    try:
        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
        if not search_docs:
            logging.info(f"search_from_wiki: No results found for '{query}'.")
            return {"wiki_results": f"No Wikipedia results found for '{query}'."}

        formatted_search_docs = "\n\n---\n\n".join(
            [
                f'<Document source="{doc.metadata.get("source", "Wikipedia")}" '
                f'page="{doc.metadata.get("page", "")}" '
                f'title="{doc.metadata.get("title", "No Title")}">\n'
                f'{doc.page_content}\n</Document>'
                for doc in search_docs
            ]
        )
        logging.info(f"search_from_wiki: Successfully fetched {len(search_docs)} results for '{query}'.")
        return {"wiki_results": formatted_search_docs}
    except Exception as e:
        logging.error(f"search_from_wiki: An error occurred for query '{query}': {e}", exc_info=True)
        return {"wiki_results": f"An error occurred during Wikipedia search for '{query}': {e}. Please try a different query or source."}

@tool
def search_from_arxiv(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 relevant paper abstracts/summaries.
    
    Args:
        query: The search query for Arxiv."""
    logging.info(f"Tool Call: search_from_arxiv(query='{query}')")
    try:
        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
        
        if not search_docs:
            logging.info(f"search_from_arxiv: No results found for '{query}'.")
            return {"arvix_results": f"No Arxiv results found for '{query}'."}

        formatted_search_docs = "\n\n---\n\n".join(
            [
                f'<Document source="{doc.metadata.get("source", "Arxiv")}" '
                f'published="{doc.metadata.get("Published", "Unknown")}" '
                f'title="{doc.metadata.get("Title", "No Title")}" '
                f'authors="{", ".join(doc.metadata.get("Authors", ["Unknown"]))}" '
                f'id="{doc.metadata.get("Id", "No ID")}">\n'
                f'{doc.page_content[:1500]}...\n</Document>'
                for doc in search_docs
            ]
        )
        logging.info(f"search_from_arxiv: Successfully fetched {len(search_docs)} results for '{query}'.")
        return {"arvix_results": formatted_search_docs}
    except Exception as e:
        logging.error(f"search_from_arxiv: An error occurred for query '{query}': {e}", exc_info=True)
        return {"arvix_results": f"An error occurred during Arxiv search for '{query}': {e}. Please try a different query or source."}



@tool
def search_from_web(query: str) -> str:
    """Search Google and return top results."""
    api_key = os.getenv("GOOGLE_API_KEY")
    cse_id = os.getenv("GOOGLE_CSE_ID")
    service = build("customsearch", "v1", developerKey=api_key)
    res = service.cse().list(q=query, cx=cse_id, num=3).execute()
    results = res.get("items", [])
    return "\n\n---\n\n".join(f"{item['title']}\n{item['link']}\n{item['snippet']}" for item in results)


# List all tools available to the agent
tools = [perform_calculation, search_from_arxiv, search_from_wiki,search_from_web]