from langchain_community.tools import DuckDuckGoSearchRun
from langchain.tools import Tool
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
from langchain_community.tools.arxiv.tool import ArxivQueryRun
from langchain_community.tools.wikipedia.tool import WikipediaQueryRun
from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
from langchain_tavily import TavilySearch
from langchain_core.tools import tool
import cmath
import pandas as pd
import numpy as np
from typing import List, Dict, Any, Optional
import tempfile
from urllib.parse import urlparse
import os
import uuid
import requests
from PIL import Image
from code_interpreter import execute_code_multilang, load_code_file
from multimodal_tools import analyse_image, analyse_audio

########################## Search Tools ##########################

@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 3 results.
    Args:
        query: The search query."""
    search_docs = WikipediaLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n</Document>{doc.page_content}\n</Document>'
            for doc in search_docs
        ]
    )
    return {"wiki_results": formatted_search_docs}


@tool
def web_search(query: str) -> str:
    """Search the web for a query using Tavily search engine and return maximum 3 results.
    Args:
        query: The search query."""
    search_docs = TavilySearch(max_results=3).invoke(input=query)
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document title="{doc["title"]}" url="{doc["url"]}" confidence score="{doc["score"]}"/>\n</Document>{doc["content"]}\n</Document>\n'
            for doc in search_docs["results"]
        ]
    )
    return {"web_results": 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, doc_content_chars_max=10000, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document title="{doc.metadata["Title"]}" page="{doc.metadata.get("page", "")}"/>\n</Document>{doc.page_content}\n</Document>\n'
            for doc in search_docs
        ]
    )
    return {"arxiv_results": formatted_search_docs}

#search_tools = [web_search,wiki_search,arxiv_search]
search_tools = [web_search,wiki_search,arxiv_search]

########################## Math tools ##########################

@tool
def add(a:float, b:float) -> float:
    """ Adds two numbers
    Args:
        a: the first number 
        b: the second number
    """

    return a+b

@tool
def subtract(a:float, b:float) -> float:
    """ Subtracts two numbers
    Args:
        a: the first number 
        b: the second number
    """

    return a-b

@tool
def multiply(a:float, b:float) -> float:
    """ Multiplies two numbers
    Args:
        a: the first number 
        b: the second number
    """

    return a*b

@tool
def divide(a:float, b:float) -> float:
    """ Divides two numbers
    Args:
        a: the first number 
        b: the second number
    """

    return a/b

@tool
def modulus(a: int, b: int) -> int:
    """
    Get the modulus of two numbers.
    Args:
        a (int): the first number
        b (int): the second number
    """
    return a % b


@tool
def power(a: float, b: float) -> float:
    """
    Get the power of two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a**b


@tool
def square_root(a: float) -> float | complex:
    """
    Get the square root of a number.
    Args:
        a (float): the number to get the square root of
    """
    if a >= 0:
        return a**0.5
    return cmath.sqrt(a)

math_tools = [add,subtract,multiply,divide,modulus,power,square_root]

########################## Doc tools ##########################

@tool
def analyze_csv_file(file_path: str) -> str:
    """
    Loads a CSV file using pandas and returns the content plus a summary of the statistics.
    Args:
        file_path (str): the path to the CSV file.
    """
    try:
        # Read the CSV file
        df = pd.read_csv(file_path)

        # Run various analyses based on the query
        result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        result += f"Columns: {', '.join(df.columns)}\n\n"

        # Add summary statistics
        result += "Summary statistics:\n"
        result += str(df.describe())

        result += "\nData:\n"
        result += f'\n{df.to_string()}\n'
        
        return result

    except Exception as e:
        return f"Error analyzing CSV file: {str(e)}"


@tool
def analyze_excel_file(file_path: str, query: str) -> str:
    """
    Loads a excel file using pandas and returns the content plus a summary of the statistics.
    Args:
        file_path (str): the path to the Excel file.
    """
    try:
        # Read the Excel file
        df = pd.read_excel(file_path)

        # Run various analyses based on the query
        result = (
            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        )
        result += f"Columns: {', '.join(df.columns)}\n\n"

        # Add summary statistics
        result += "Summary statistics:\n"
        result += str(df.describe())

        result += "\nData:\n"
        result += f'\n{df.to_string()}\n'

        return result

    except Exception as e:
        return f"Error analyzing Excel file: {str(e)}"


doc_tools = [analyze_csv_file,analyze_excel_file]

######################### Code tools #########################

code_tools = [execute_code_multilang,load_code_file]

######################### Image tools #########################

multimodal_tools = [analyse_image, analyse_audio]

######################### Final tools #########################

def get_tools():
    return search_tools + multimodal_tools + code_tools + doc_tools + math_tools