from langchain_community.tools import DuckDuckGoSearchRun
import base64
from langchain.tools import Tool
import random
from langchain.tools import Tool
from huggingface_hub import list_models
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
from langchain_community.document_loaders.csv_loader import CSVLoader
import math
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
import os
from typing import Optional, Tuple
from huggingface_hub import list_models
from pathlib import Path
import json 
from langchain_community.document_loaders import UnstructuredExcelLoader
from langchain_tavily import TavilySearch
from langchain.schema import Document
from langchain_community.document_loaders import YoutubeLoader
from langchain_community.document_loaders import AssemblyAIAudioTranscriptLoader
from langchain_community.document_loaders.assemblyai import TranscriptFormat
import assemblyai as aai
import tempfile
from typing import Dict, List
from langchain_community.document_loaders import GoogleApiClient, GoogleApiYoutubeLoader
from langchain_openai import ChatOpenAI

def create_credentials_file(api_key: str) -> str:
    """
    Create a temporary credentials file for Google API access.
    
    This function creates the necessary credentials structure that Google's API expects.
    In a production environment, you would typically have a proper service account
    JSON file, but for this educational example, we'll create a minimal structure.
    
    Args:
        api_key: "The API Key"
        
    Returns:
        Path to the temporary credentials file
    """
    # Create a minimal credentials structure
    # Note: In production, you would use proper service account credentials
    credentials_data = {
        "type": "service_account",
        "project_id": "your-project-id",
        "private_key_id": "dummy-key-id",
        "private_key": "dummy-private-key",
        "client_email": "tu101@pnw.edu",
        "client_id": "dummy-client-id",
        "auth_uri": "https://accounts.google.com/o/oauth2/auth",
        "token_uri": "https://oauth2.googleapis.com/token",
        "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
        "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/dummy%40your-project-id.iam.gserviceaccount.com"
    }
    
    # Create a temporary file to store credentials
    temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
    json.dump(credentials_data, temp_file)
    temp_file.close()
    
    return temp_file.name

def extract_video_id_from_url(url: str) -> Optional[str]:
    """
    Extract the video ID from a YouTube URL.
    
    This function handles different YouTube URL formats and extracts the unique
    video identifier that the API needs to fetch the video content.
    
    Args:
        url: YouTube URL in various formats
        
    Returns:
        Video ID string or None if URL is invalid
    """
    import re
    
    # Handle different YouTube URL formats
    patterns = [
        r'(?:https?://)?(?:www\.)?youtube\.com/watch\?v=([a-zA-Z0-9_-]+)',
        r'(?:https?://)?(?:www\.)?youtu\.be/([a-zA-Z0-9_-]+)',
        r'(?:https?://)?(?:www\.)?youtube\.com/embed/([a-zA-Z0-9_-]+)',
    ]
    
    for pattern in patterns:
        match = re.search(pattern, url)
        if match:
            return match.group(1)
    
    return None

def youtube_search_google_api(link: str) -> Dict:
    """
    Search YouTube using Google's official API and extract transcript/captions.
    
    This function represents a more robust approach to YouTube content access.
    Instead of web scraping, we use Google's official API endpoints, which
    provide reliable access to video metadata and transcript data.
    
    Args:
        link: YouTube URL to process
        
    Returns:
        Dictionary containing video documents or error information
    """
    
    # Your Google Cloud API key
    api_key = "AIzaSyA-1ZuYz-PYnnE9PQsWA7vOuJeRz-E2sbU"
    
    try:
        # First, let's validate and extract the video ID
        video_id = extract_video_id_from_url(link)
        if not video_id:
            return {
                "YT_docs": [{
                    "text": "Invalid YouTube URL format. Please provide a valid YouTube link.",
                    "metadata": {"error": True, "url": link}
                }]
            }
        
        # For this implementation, we'll use a simpler approach with the API key
        # In a full production system, you would use proper service account credentials
        
        # Alternative approach using youtube-transcript-api for transcript data
        # This is often more reliable for transcript access
        try:
            from youtube_transcript_api import YouTubeTranscriptApi
            
            # Get the transcript for the video
            transcript_list = YouTubeTranscriptApi.get_transcript(video_id)
            
            # Combine all transcript segments into a single text
            full_transcript = ""
            for segment in transcript_list:
                full_transcript += segment['text'] + " "
            
            # Create a document-like structure similar to other loaders
            documents = [{
                "text": full_transcript.strip(),
                "metadata": {
                    "video_id": video_id,
                    "source": link,
                    "title": f"YouTube Video {video_id}",
                    "segments_count": len(transcript_list)
                }
            }]
            
            return {"YT_docs": documents}
            
        except Exception as transcript_error:
            # If transcript API fails, provide helpful error information
            error_msg = f"Could not retrieve transcript for video {video_id}. "
            error_msg += "This video may not have captions available, may be private, "
            error_msg += "or may have restricted access. "
            error_msg += f"Original error: {str(transcript_error)}"
            
            return {
                "YT_docs": [{
                    "text": error_msg,
                    "metadata": {"error": True, "video_id": video_id}
                }]
            }
    
    except Exception as e:
        # Comprehensive error handling with debugging information
        error_msg = f"Error processing YouTube video: {str(e)}"
        return {
            "YT_docs": [{
                "text": error_msg,
                "metadata": {"error": True, "url": link}
            }]
        }

YT_search_tool = Tool(
    name="YT_search",
    func=youtube_search_google_api,
    description="Searches YouTube using Google's official services to extract video transcripts and captions. Handles various YouTube URL formats and provides detailed error information when content is not available."
)


def audio_search(path: str) -> str:
    """Read an audio mp3 file and extract the text to understand it. 
    
    Args:
        path: The path to the file. If the file attached is just {name}.mp3, then it is local, and just use {name}.mp3. """
    aai.settings.api_key = "945f7d6c00f142158322cd0c13901024"
    possible_paths = [
        path,  # Try the path as-is first
        os.path.join(".", path),  # Try current directory
        os.path.join("/tmp", path),  # Try common temp directory
        os.path.join("uploads", path),  # Try uploads directory
    ]
    
    # Find the actual file location
    actual_path = None
    for potential_path in possible_paths:
        if os.path.exists(potential_path):
            actual_path = potential_path
            break
    
    if actual_path is None:
        # If we can't find the file, let's provide debugging information
        current_dir_files = os.listdir(".")
        error_msg = f"Could not find audio file. Searched in: {possible_paths}. Current directory contains: {current_dir_files[:10]}"  # Limit output for readability
        return {"audio_docs": [{"text": error_msg, "metadata": {"error": True}}]}
    
    try:
        # Use the found path
        loader = AssemblyAIAudioTranscriptLoader(
            file_path=actual_path, 
            transcript_format=TranscriptFormat.PARAGRAPHS
        )
        
        raw_docs = loader.load()
        
        if not raw_docs:
            return {"audio_docs": [{"text": "No audio content could be transcribed from this file.", "metadata": {"warning": True}}]}
        
        # Convert each Document into a serializable dict
        serialized_docs = [
            {"text": doc.page_content, "metadata": doc.metadata}
            for doc in raw_docs
        ]
        
        return {"audio_docs": serialized_docs}
        
    except Exception as e:
        error_msg = f"Error processing audio file at {actual_path}: {str(e)}"
        return {"audio_docs": [{"text": error_msg, "metadata": {"error": True}}]}

        


audio_search_tool = Tool(
    name="audio_search",
    func=audio_search,
    description="Uses an mp3 audio file path and extract it into paragraphs of text that are searchable."
)


def excel_search(excel_path: str ) -> str:
    """Given the path to an excel file, return the excel file contents as raw text.
    
    Args:
        excel_path: The path to where the excel file is stored."""
    loader = UnstructuredExcelLoader(excel_path, mode="elements")
    raw_excel_docs = loader.load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in raw_excel_docs
        ])
    return {"excel_results": formatted_search_docs}

excel_search_tool = Tool(
    name="excel_search",
    func=excel_search,
    description="Searches for and retrieves information in an Excel Document using the given path."
)
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 source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return {"wiki_results": formatted_search_docs}

wiki_search_tool = Tool(
    name="wiki_search",
    func=wiki_search,
    description="Searches for and retrieves information on Wikipedia using an inputted query."
)

def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 5 results. Make sure to format the args as query and have the query sound like a sentence, hopefully matching part of the question the search is used for.
    Args:
    query: The search query. The query fulfils the argument that is input."""
    tavily = TavilySearch(max_results=5, topic="general", include_answer=True, include_image_descriptions=True)
    raw_results = tavily.invoke(query)
    
    # Extract the actual search results from the response dictionary
    # The 'results' key contains the list of search results we want
    search_results = raw_results.get('results', [])
    
    docs = []
    for idx, result in enumerate(search_results):
        # Each result is a dictionary with keys like 'title', 'content', 'url'
        # We'll use the content as the main text and include other metadata
        content = result.get('content', result.get('title', 'No content available'))
        source_url = result.get('url', f'Tavily result #{idx+1}')
        
        docs.append(Document(
            page_content=content,
            metadata={"source": source_url, "title": result.get('title', 'No title')}
        ))
    
    formatted_docs = "\n\n---\n\n".join(
        f'<Document source="{doc.metadata["source"]}">\n'
        f'{doc.page_content}\n'
        f'</Document>'
        for doc in docs
    )
    
    # Fix the return type - should return string, not dictionary
    return formatted_docs

web_search_tool = Tool(
    name="web_search",
    func=web_search,
    description="Searches for and retrieves information online using an inputted query using Tavily. Please use a query that sounds like a sentence and not just keywords."
)

def arvix_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 results.
    
    Args:
        query: The search query.
    
    Returns:
        Formatted string containing search results from Arxiv papers.
    """
    try:
        # Load documents from Arxiv
        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
        
        # Check if we got any results
        if not search_docs:
            return "No results found for the query: " + query
        
        # Process each document with safe metadata access
        formatted_docs = []
        for doc in search_docs:
            # Use .get() method for safe dictionary access with fallback values
            source = doc.metadata.get("source", "Unknown source")
            page = doc.metadata.get("page", "N/A")
            
            # Create the formatted document string
            formatted_doc = f'<Document source="{source}" page="{page}"/>\n{doc.page_content[:1000]}\n</Document>'
            formatted_docs.append(formatted_doc)
        
        # Join all documents with separators
        formatted_search_docs = "\n\n---\n\n".join(formatted_docs)
        
        # Return the string directly (not wrapped in a dictionary)
        return formatted_search_docs
        
    except Exception as e:
        # Handle any unexpected errors gracefully
        return f"Error searching Arxiv: {str(e)}"
        
arvix_search_tool = Tool(
    name="arvix_search",
    func=arvix_search,
    description="Searches for and retrieves information from Arxiv, a database storing scientific papers. Returns formatted text containing paper abstracts and metadata."
)

def extract_image_text(img_path: str) -> str:
    """
    Extract text from an image file using a multimodal model.

    Args:
        img_path: A local image file path (strings).

    Returns:
        A single string containing the concatenated text extracted from each image.
    """
    vision_llm = ChatOpenAI(model="gpt-4o")
    all_text = ""
    try:

        # Read image and encode as base64
        with open(img_path, "rb") as image_file:
            image_bytes = image_file.read()

        image_base64 = base64.b64encode(image_bytes).decode("utf-8")

        # Prepare the prompt including the base64 image data
        message = [
            HumanMessage(
                content=[
                    {
                        "type": "text",
                        "text": (
                            "Extract all the text from this image. "
                            "Return only the extracted text, no explanations."
                        ),
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{image_base64}"
                        },
                    },
                ]
            )
        ]

        # Call the vision-capable model
        response = vision_llm.invoke(message)

        # Append extracted text
        all_text += response.content + "\n\n"

        return all_text.strip()
    except Exception as e:
        # You can choose whether to raise or just return an empty string / error message
        error_msg = f"Error extracting text: {str(e)}"
        print(error_msg)
        return ""
    
extract_image_text = Tool(
    name="extract_image_text",
    func=extract_image_text,
    description="Extracts text that might be visible on an image by reading the image. Works with path/url to image."
)


def multiply(a: int, b: int) -> int:
    """Multiply two numbers.
    Args:
        a: first int
        b: second int
    """
    return a * b

multiply = Tool(
    name="multiply",
    func=multiply,
    description="Multiplies two integer numbers to get the product."
)

def add(a: int, b: int) -> int:
    """Add two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a + b
    
add = Tool(
    name="add",
    func=add,
    description="Adds two integer numbers to get the result."
)

def subtract(a: int, b: int) -> int:
    """Subtract two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a - b
subtract = Tool(
    name="subtract",
    func=subtract,
    description="Subtracts two integer numbers to get the result."
)

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
divide = Tool(
    name="divide",
    func=divide,
    description="Divides two integer numbers to get the result."
)

def modulus(a: int, b: int) -> int:
    """Get the modulus of two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a % b
    
modulus = Tool(
    name="modulus",
    func=modulus,
    description="Calculates the modulus of two integer numbers, a mod b, to get the result."
)

def calculate_cargo_travel_time(
    origin_coords: Tuple[float, float],
    destination_coords: Tuple[float, float],
    cruising_speed_kmh: Optional[float] = 750.0,  # Average speed for cargo planes
) -> float:
    """
    Calculate the travel time for an object, usually a cargo plane, between two points on Earth using great-circle distance.

    Args:
        origin_coords: Tuple of (latitude, longitude) for the starting point
        destination_coords: Tuple of (latitude, longitude) for the destination
        cruising_speed_kmh: Optional cruising speed in km/h (defaults to 750 km/h for typical cargo planes). Adjustable depending on object.

    Returns:
        float: The estimated travel time in hours

    Example:
        >>> # Chicago (41.8781° N, 87.6298° W) to Sydney (33.8688° S, 151.2093° E)
        >>> result = calculate_cargo_travel_time((41.8781, -87.6298), (-33.8688, 151.2093))
    """

    def to_radians(degrees: float) -> float:
        return degrees * (math.pi / 180)

    # Extract coordinates
    lat1, lon1 = map(to_radians, origin_coords)
    lat2, lon2 = map(to_radians, destination_coords)

    # Earth's radius in kilometers
    EARTH_RADIUS_KM = 6371.0

    # Calculate great-circle distance using the haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1

    a = (
        math.sin(dlat / 2) ** 2
        + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2
    )
    c = 2 * math.asin(math.sqrt(a))
    distance = EARTH_RADIUS_KM * c

    # Add 10% to account for non-direct routes and air traffic controls
    actual_distance = distance * 1.1

    # Calculate flight time
    # Add 1 hour for takeoff and landing procedures
    flight_time = (actual_distance / cruising_speed_kmh) + 1.0

    # Format the results

    return round(flight_time, 2)

calculate_cargo_travel_time = Tool(
    name="calculate_cargo_travel_time",
    func=calculate_cargo_travel_time,
    description="Calculates the travel time between two points on the earth based on an object with optionally set speed, usually a cargo plane."
)


def get_hub_stats(author: str) -> str:
    """Fetches the most downloaded model from a specific author on the Hugging Face Hub."""
    try:
        # List models from the specified author, sorted by downloads
        models = list(list_models(author=author, sort="downloads", direction=-1, limit=1))

        if models:
            model = models[0]
            return f"The most downloaded model by {author} is {model.id} with {model.downloads:,} downloads."
        else:
            return f"No models found for author {author}."
    except Exception as e:
        return f"Error fetching models for {author}: {str(e)}"

# Initialize the tool
hub_stats_tool = Tool(
    name="get_hub_stats",
    func=get_hub_stats,
    description="Fetches the most downloaded model from a specific author on the Hugging Face Hub."
)