from langchain.tools import Tool
import requests
import os
from PIL import Image
import io
import base64
from ddgs import DDGS
from typing import Optional
import json
import PyPDF2
import tempfile
import requests
from bs4 import BeautifulSoup
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.schema import Document
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

def web_search_tool_func(query: str) -> str:
    """Enhanced web search with Wikipedia priority using DDGS."""
    try:
        # Try Wikipedia-specific search first
        print(f"🔍 Performing web search for: {query}")
        wiki_results = search_wikipedia(query)
        if wiki_results and len(wiki_results.strip()) > 100:  # Good Wikipedia result
            return f"Wikipedia search results:\n{wiki_results}"
        
        # Fall back to general web search
        general_results = search_general(query)
        if general_results:
            return f"Web search results:\n{general_results}"
        else:
            return "No relevant search results found."
            
    except Exception as e:
        return f"Web search failed: {str(e)}"

def search_wikipedia(query: str) -> str:
    """Search Wikipedia specifically for factual information."""
    try:
        with DDGS() as ddgs:
            # Try multiple Wikipedia search strategies
            search_queries = [
                f"site:en.wikipedia.org {query}",  # English Wikipedia specifically
                f"{query} site:wikipedia.org",     # Alternative format
                f"{query} wikipedia"               # General Wikipedia search
            ]
            
            for search_query in search_queries:
                try:
                    results = list(ddgs.text(search_query, max_results=3))
                    
                    if results:
                        # Filter for relevant Wikipedia results
                        wiki_results = []
                        for result in results:
                            title = result.get('title', 'No title')
                            body = result.get('body', 'No content')
                            url = result.get('href', '')
                            
                            # Only include if it's actually Wikipedia and relevant
                            if 'wikipedia.org' in url.lower() and any(term in title.lower() or term in body.lower() for term in query.lower().split()):
                                wiki_results.append(f"Title: {title}\nContent: {body}\nSource: {url}\n")
                        
                        if wiki_results:
                            return "\n---\n".join(wiki_results)
                
                except Exception:
                    continue  # Try next search query
            
            return ""  # No good results found
            
    except Exception as e:
        return f"Wikipedia search failed: {str(e)}"

def search_general(query: str) -> str:
    """General web search as fallback."""
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=5))
            
            if not results:
                return ""
            
            # Format general results
            formatted_results = []
            for result in results:
                title = result.get('title', 'No title')
                body = result.get('body', 'No content')
                url = result.get('href', '')
                
                # Prioritize reliable sources
                if any(domain in url.lower() for domain in ['wikipedia.org', 'britannica.com', 'edu', 'gov']):
                    formatted_results.insert(0, f"Title: {title}\nContent: {body}\nSource: {url}\n")
                else:
                    formatted_results.append(f"Title: {title}\nContent: {body}\nSource: {url}\n")
            
            return "\n---\n".join(formatted_results)
            
    except Exception as e:
        return f"General search failed: {str(e)}"

web_search_tool = Tool(
    name="web_search",
    func=web_search_tool_func,
    description="Searches the web for current information. Use this for factual questions, recent events, or when you need to find information not in your training data."
)

def file_download_tool_func(task_id: str) -> str:
    """Downloads a file associated with a GAIA task ID."""
    try:
        api_url = "https://agents-course-unit4-scoring.hf.space"
        file_url = f"{api_url}/files/{task_id}"
        
        response = requests.get(file_url, timeout=30)
        response.raise_for_status()
        
        # Save to temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".tmp") as temp_file:
            temp_file.write(response.content)
            temp_path = temp_file.name
        
        # Try to determine file type and process accordingly
        content_type = response.headers.get('content-type', '').lower()
        
        if 'image' in content_type:
            return f"Image file downloaded to {temp_path}. Use image_analysis_tool to analyze it."
        elif 'pdf' in content_type:
            return process_pdf_file(temp_path)
        elif 'text' in content_type:
            with open(temp_path, 'r', encoding='utf-8') as f:
                content = f.read()
            os.unlink(temp_path)  # Clean up
            return f"Text file content:\n{content}"
        else:
            return f"File downloaded to {temp_path}. Content type: {content_type}"
            
    except Exception as e:
        return f"Failed to download file for task {task_id}: {str(e)}"

def process_pdf_file(file_path: str) -> str:
    """Process a PDF file and extract text content."""
    try:
        with open(file_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            text_content = ""
            
            for page_num in range(len(pdf_reader.pages)):
                page = pdf_reader.pages[page_num]
                text_content += f"\n--- Page {page_num + 1} ---\n"
                text_content += page.extract_text()
        
        os.unlink(file_path)  # Clean up
        return f"PDF content extracted:\n{text_content}"
    except Exception as e:
        return f"Failed to process PDF: {str(e)}"

file_download_tool = Tool(
    name="file_download",
    func=file_download_tool_func,
    description="Downloads and processes files associated with GAIA task IDs. Can handle images, PDFs, and text files."
)

def image_analysis_tool_func(image_path_or_description: str) -> str:
    """Analyzes images for GAIA questions. For now, returns a placeholder."""
    # This is a simplified version - in a full implementation, you'd use a vision model
    try:
        if os.path.exists(image_path_or_description):
            # Try to open and get basic info about the image
            with Image.open(image_path_or_description) as img:
                width, height = img.size
                mode = img.mode
                format_info = img.format
                
                # Clean up the temporary file
                os.unlink(image_path_or_description)
                
                return f"Image analyzed: {width}x{height} pixels, mode: {mode}, format: {format_info}. Note: This is a basic analysis. For detailed image content analysis, a vision model would be needed."
        else:
            return f"Image analysis requested for: {image_path_or_description}. Note: Full image analysis requires a vision model integration."
    except Exception as e:
        return f"Image analysis failed: {str(e)}"

image_analysis_tool = Tool(
    name="image_analysis",
    func=image_analysis_tool_func,
    description="Analyzes images to extract information. Use this for questions involving visual content."
)

def calculator_tool_func(expression: str) -> str:
    """Performs mathematical calculations safely."""
    try:
        # Basic safety check - only allow certain characters
        allowed_chars = set('0123456789+-*/().= ')
        if not all(c in allowed_chars for c in expression):
            return f"Invalid characters in expression: {expression}"
        
        # Use eval safely for basic math
        result = eval(expression)
        return f"Calculation result: {expression} = {result}"
    except Exception as e:
        return f"Calculation failed for '{expression}': {str(e)}"

calculator_tool = Tool(
    name="calculator",
    func=calculator_tool_func,
    description="Performs mathematical calculations. Use this for numerical computations and math problems."
)

def text_processor_tool_func(text: str, operation: str = "summarize") -> str:
    """Processes text for various operations like summarization, extraction, etc."""
    try:
        if operation == "summarize":
            # Simple summarization - take first and last sentences if long
            sentences = text.split('.')
            if len(sentences) > 5:
                summary = '. '.join(sentences[:2] + sentences[-2:])
                return f"Text summary: {summary}"
            else:
                return f"Text (short enough to not need summarization): {text}"
        
        elif operation == "extract_numbers":
            import re
            numbers = re.findall(r'\d+(?:\.\d+)?', text)
            return f"Numbers found in text: {numbers}"
        
        elif operation == "extract_dates":
            import re
            # Simple date pattern matching
            date_patterns = [
                r'\d{1,2}/\d{1,2}/\d{4}',  # MM/DD/YYYY
                r'\d{4}-\d{1,2}-\d{1,2}',  # YYYY-MM-DD
                r'\b\w+ \d{1,2}, \d{4}\b'  # Month DD, YYYY
            ]
            dates = []
            for pattern in date_patterns:
                dates.extend(re.findall(pattern, text))
            return f"Dates found in text: {dates}"
        
        else:
            return f"Text processing operation '{operation}' not supported. Available: summarize, extract_numbers, extract_dates"
            
    except Exception as e:
        return f"Text processing failed: {str(e)}"

text_processor_tool = Tool(
    name="text_processor",
    func=text_processor_tool_func,
    description="Processes text for various operations like summarization, number extraction, date extraction. Specify operation as second parameter."
)

def enhanced_web_retrieval_tool_func(query: str) -> str:
    """Enhanced web search with vector retrieval for deep content analysis."""
    try:
        print(f"🔍 Enhanced web retrieval for: {query}")
        
        # Step 1: Get search results with URLs
        search_results = get_search_urls(query)
        if not search_results:
            return "No search results found."
        
        # Step 2: Fetch and process webpage content
        documents = []
        for result in search_results[:4]:  # Top 4 results as requested
            url = result.get('url', '')
            title = result.get('title', 'No title')
            
            print(f"📄 Fetching content from: {title}")
            content = fetch_webpage_content(url)
            if content:
                doc = Document(
                    page_content=content,
                    metadata={"source": url, "title": title}
                )
                documents.append(doc)
        
        if not documents:
            return "Could not fetch content from any search results."
        
        # Step 3: Create vector store and search
        return search_documents_with_vector_store(documents, query)
        
    except Exception as e:
        return f"Enhanced web retrieval failed: {str(e)}"

def get_search_urls(query: str) -> list:
    """Get search results from English Wikipedia only using DDGS."""
    try:
        with DDGS() as ddgs:
            # Create Wikipedia-specific search queries
            wikipedia_queries = [
                f"site:en.wikipedia.org {query}",
                f"{query} site:en.wikipedia.org"
            ]
            
            search_results = []
            seen_urls = set()
            
            for wiki_query in wikipedia_queries:
                try:
                    results = list(ddgs.text(wiki_query, max_results=2))
                    
                    for result in results:
                        url = result.get('href', '')
                        
                        # Only include Wikipedia URLs and avoid duplicates
                        if 'en.wikipedia.org' in url and url not in seen_urls:
                            search_results.append({
                                'url': url,
                                'title': result.get('title', 'No title'),
                                'snippet': result.get('body', 'No content')
                            })
                            seen_urls.add(url)
                            
                            # Limit to 4 unique Wikipedia pages
                            if len(search_results) >= 4:
                                break
                    
                    if len(search_results) >= 4:
                        break
                        
                except Exception:
                    continue  # Try next query
            
            return search_results
            
    except Exception as e:
        print(f"Wikipedia search URL retrieval failed: {e}")
        return []

def fetch_webpage_content(url: str) -> str:
    """Fetch and extract clean text content from a webpage."""
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        
        response = requests.get(url, headers=headers, timeout=10)
        response.raise_for_status()
        
        # Parse HTML and extract text
        soup = BeautifulSoup(response.content, 'html.parser')
        
        # Remove script and style elements
        for script in soup(["script", "style"]):
            script.decompose()
        
        # Get text content
        text = soup.get_text()
        
        # Clean up text
        lines = (line.strip() for line in text.splitlines())
        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
        text = ' '.join(chunk for chunk in chunks if chunk)
        
        return text[:20000]  # Increase to 20k characters to get more content
        
    except Exception as e:
        print(f"Failed to fetch content from {url}: {e}")
        return ""

def search_documents_with_vector_store(documents: list, query: str) -> str:
    """Create vector store and search for relevant information."""
    try:
        # Split documents into chunks
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
        )
        
        splits = text_splitter.split_documents(documents)
        
        if not splits:
            return "No content to process after splitting."
        
        # Create embeddings and vector store
        embeddings = OpenAIEmbeddings()
        vectorstore = FAISS.from_documents(splits, embeddings)
        
        # Search for relevant chunks with the original query
        relevant_docs = vectorstore.similarity_search(query, k=5)
        
        # Format results
        results = []
        for i, doc in enumerate(relevant_docs, 1):
            source = doc.metadata.get('source', 'Unknown source')
            title = doc.metadata.get('title', 'No title')
            content = doc.page_content[:5000]  # First 500 chars
            
            results.append(f"Result {i} from {title}:\n{content}\nSource: {source}\n")
        
        return "\n---\n".join(results)
        
    except Exception as e:
        return f"Vector search failed: {str(e)}"

enhanced_web_retrieval_tool = Tool(
    name="enhanced_web_retrieval",
    func=enhanced_web_retrieval_tool_func,
    description="Enhanced Wikipedia-only search with vector retrieval. Fetches full content from English Wikipedia pages and uses semantic search to find relevant information. Use this for factual questions that need detailed Wikipedia content analysis."
)

# List of all tools for easy import
agent_tools = [
    web_search_tool,
    enhanced_web_retrieval_tool,
    file_download_tool,
    image_analysis_tool,
    calculator_tool,
    text_processor_tool
]