import requests
import pandas as pd 
from io import BytesIO
from typing import List 
from langchain_core.tools import tool
from langchain_community.tools import DuckDuckGoSearchResults
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import ArxivLoader, WikipediaLoader

@tool
def add(a: int, b: int) -> int:
    """Add two numbers.

    Args:
        a: first int
        b: second int
    """
    return a + b


@tool
def subtract(a: int, b: int) -> int:
    """Subtract two numbers.

    Args:
        a: first int
        b: second int
    """
    return a - b


@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers.

    Args:
        a: first int
        b: second int
    """
    return a * b


@tool
def divide(a: int, b: int) -> int:
    """Divide two numbers.

    Args:
        a: first int
        b: second int
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b


@tool
def modulo(a: int, b: int) -> int:
    """Get the modulus of two numbers.

    Args:
        a: first int
        b: second int
    """
    return a % b

@tool
def add_list(lst: List[float]) -> float:
    """Sum up a list of numbers.

    Args:
        lst: A list of numbers.
    """
    return sum(lst)

@tool
def web_search(query: str) -> str:
    """Web search with Tavily for a query and return maximum 3 results.

    Args:
        query: The search query."""
    # search_docs = DuckDuckGoSearchResults(max_results=3).invoke({'query': query})
    search_docs = TavilySearchResults(max_results=3).invoke({'query': query})
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document title="{doc["title"]}>"\n{doc["content"]}\n</Document>'
            for doc in search_docs
        ])
    return f"web_results:\n{formatted_search_docs}"


@tool
def arxiv_search(query: str) -> str:
    """Search ArXiv for a query and return maximum 3 result.

    Args:
        query: The search query."""
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document title="{doc.metadata["Title"]}>"\n{doc.metadata["Summary"]}\n</Document>'
            for doc in search_docs
        ])
    return f"arxiv_content:\n{formatted_search_docs}"


@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 results.

    Args:
        query: The search query."""
    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document title="{doc.metadata["Title"]}>"\n{doc.metadata["Summary"]}\n</Document>'
            for doc in search_docs
        ])
    return f"wikipedia_content:\n{formatted_search_docs}"


@tool
def read_xlsx_file(file_location: str) -> str:
    """Read an .xlsx file from a path and return its contents as a string.
    
    Args:
        file_location (str): Path to the xlsx file.
    """
    base_name = file_location.split(".")[0]
    file_url = f"https://agents-course-unit4-scoring.hf.space/files/{base_name}"

    # Download the file
    response = requests.get(file_url)
    if response.status_code != 200:
        raise RuntimeError(f"Failed to download file: {file_url}")

    # Load the xlsx file into a pandas ExcelFile object
    df_dict = pd.read_excel(BytesIO(response.content), sheet_name=None)

    # Convert all sheets to a single formatted string
    text_output = ""
    for sheet_name, sheet_df in df_dict.items():
        text_output += f"--- Sheet: {sheet_name} ---\n"
        text_output += sheet_df.to_string(index=False)
        text_output += "\n\n"

    return f"xlsx file content:\n{text_output.strip()}"

@tool
def get_python_file(file_location: str) -> str:
    """
    Get the content of a Python file from a path.
    Args:
        file_location: path to the Python file.
    """
    base_name = file_location.split(".")[0]
    file_url = f"https://agents-course-unit4-scoring.hf.space/files/{base_name}"

    # Download the file
    response = requests.get(file_url)
    if response.status_code != 200:
        raise RuntimeError(f"Failed to download file: {file_url}")

    return f'python file:\n{response.content}'