from langchain.tools import tool
#from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.tools import TavilySearchResults
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.document_loaders import YoutubeLoader
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
import requests
import os
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI

@tool
def add(a: float, b: float) -> float:
    """Add two integers and return the result."""
    return a + b

@tool
def subtract(a: float, b: float) -> float:
    """Subtract two integers and return the result."""
    return a - b

@tool
def multiply(a: float, b: float) -> float:
    """Multiply two integers and return the result."""
    return a * b

@tool
def divide(a: float, b: float) -> float:
    """Divide two integers and return the result."""
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def exponentiate(base: float, exponent: float) -> float:
    """Raise a number to the power of another number and return the result."""
    return base ** exponent

@tool
def modulus(a: float, b: float) -> float:
    """Return the modulus of two integers."""
    return a % b

@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia and returns only 2 results. 
    
    Args:
        query: The search query."""
    docs = WikipediaLoader(query=query, load_max_docs=2).load()
    res = "\n#######\n".join(
        [
            f"Document {i+1}:\nSource: {doc.metadata.get('source', '')}\nPage: {doc.metadata.get('page', '')}\nContent:\n{doc.page_content}\n"
            for i, doc in enumerate(docs)
        ])
    print(f"load wiki page : {res}")
    return {"results": res}

@tool
def load_web_page(url: str) -> str:
    """Load a web page and return its content.
    
    Args:
        url: The URL of the web page to load.
    """
    loader = WebBaseLoader(url)
    docs = loader.load()
    res = "\n#######\n".join(
        [
            f"Document {i+1}:\nSource: {doc.metadata.get('source', '')}\nPage: {doc.metadata.get('page', '')}\nContent:\n{doc.page_content}\n"
            for i, doc in enumerate(docs)
        ])
    print(f"load web page : {res}")
    return {"results": res}
    
@tool
def paper_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 result.
    
    Args:
        query: The search query."""
    docs = ArxivLoader(query=query, load_max_docs=3).load()
    res = "\n#######\n".join(
        [
            f"Document {i+1}:\nSource: {doc.metadata.get('source', '')}\nPage: {doc.metadata.get('page', '')}\nContent:\n{doc.page_content}\n"
            for i, doc in enumerate(docs)
        ])
    print(f"load paper page : {res}")
    return {"results": res}

@tool
def understand_image(text: str, image_url: str):
    """
    Sends a text prompt and an image URL to OpenAI's API using the ChatOpenAI model.
    Returns the model's response.

    Args:
        text (str): The text prompt to send.
        image_url (str): URL to the image to send.

    Returns:
        str: The response from the model.
    """

    # Fetch image from URL and encode as base64
    #response = requests.get(image_url)
    #image_bytes = response.content
    #image_b64 = base64.b64encode(image_bytes).decode("utf-8")

    # Prepare message with text and image
    message = HumanMessage(
        content=[
            {"type": "text", "text": text},
            #{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}", "detail": "auto"}}
            {"type": "image_url", "image_url": {"url": image_url}}
        ]
    )
    model = ChatOpenAI(model="gpt-4o", temperature=0)
    response = model.invoke([message])
    return response.content

@tool
def load_youtube_video(url: str) -> str:
    """Load a YouTube video and return its content."""
    loader = YoutubeLoader.from_youtube_url(url, add_video_info=True)
    documents = loader.load()
    return documents[0].page_content if documents else "No content found"

@tool
def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 5 results.
    
    Args:
        query: The search query."""
    documents = TavilySearchResults(max_results=5).invoke(input=query)
    res = "\n#######\n".join(
        [
            f"Document {i+1}:\nContent: {doc['content']}\n"
            for i, doc in enumerate(documents)
        ])
    print(f"load tavily search : {res}")
    return {"results": res}

@tool
def transcribe_audio(audio_url: str) -> str:
    """Transcribe audio from a URL and return the text.
    
    Args:
        audio_url: The URL of the audio file to transcribe.
    """
    
    response = requests.get(audio_url)
    audio_file = "audio.mp3"
    with open(audio_file, "wb") as f:
        f.write(response.content)

    # Step 2: Send it to OpenAI's transcription API
    api_key = os.environ.get("OPENAI_API_KEY")
    headers = {
        "Authorization": f"Bearer {api_key}"
    }
    files = {
        'file': (audio_file, open(audio_file, 'rb')),
        'model': (None, 'whisper-1')
    }

    transcribe_response = requests.post(
        "https://api.openai.com/v1/audio/transcriptions",
        headers=headers,
        files=files
    )
    print(f"Transcription response: {transcribe_response.json()}")
    return {"results": transcribe_response.json().get("text", "Transcription failed.")}