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 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 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, backend: str = "bing") -> str:
    """Enhanced web search with cascading fallback: Wikipedia first, then general web search."""
    try:
        print(f"🔍 Enhanced web retrieval for: {query}")
        
        # Step 1: Try Wikipedia search first
        print("📚 Searching Wikipedia...")
        wikipedia_results = get_wikipedia_search_urls(query, backend)
        
        if has_sufficient_results(wikipedia_results):
            print(f"✅ Found {len(wikipedia_results)} Wikipedia results")
            documents = fetch_and_process_results(wikipedia_results, "Wikipedia")
            if documents:
                return search_documents_with_vector_store(documents, query, "Wikipedia")
        
        # Step 2: Fallback to general web search
        print("🌐 Wikipedia results insufficient, searching general web...")
        web_results = get_general_web_search_urls(query, backend)
        
        if web_results:
            print(f"✅ Found {len(web_results)} general web results")
            documents = fetch_and_process_results(web_results, "General Web")
            if documents:
                return search_documents_with_vector_store(documents, query, "General Web")
        
        return "No sufficient results found in Wikipedia or general web search."
        
    except Exception as e:
        return f"Enhanced web retrieval failed: {str(e)}"

def get_wikipedia_search_urls(query: str, backend: str = "auto") -> list:
    """Get search results from English Wikipedia using DDGS."""
    try:
        with DDGS() as ddgs:
            # Create Wikipedia-specific search queries
            wikipedia_queries = [
                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=8, 
                        region="us-en", 
                        backend=backend, 
                        safesearch="moderate"
                    ))
                    
                    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 6 unique Wikipedia pages
                            if len(search_results) >= 6:
                                break
                    
                    if len(search_results) >= 6:
                        break
                        
                except Exception as e:
                    print(f"Wikipedia search attempt failed: {e}")
                    continue
            
            return search_results
            
    except Exception as e:
        print(f"Wikipedia search URL retrieval failed: {e}")
        return []

def get_general_web_search_urls(query: str, backend: str = "auto") -> list:
    """Get search results from general web using DDGS."""
    try:
        with DDGS() as ddgs:
            search_results = []
            seen_urls = set()
            
            try:
                # General web search without site restriction
                results = list(ddgs.text(
                    query, 
                    max_results=8, 
                    region="us-en", 
                    backend=backend, 
                    safesearch="moderate"
                ))
                
                for result in results:
                    url = result.get('href', '')
                    
                    # Avoid duplicates and filter out low-quality sources
                    if url not in seen_urls and is_quality_source(url):
                        search_results.append({
                            'url': url,
                            'title': result.get('title', 'No title'),
                            'snippet': result.get('body', 'No content')
                        })
                        seen_urls.add(url)
                        
                        # Limit to 6 unique web pages
                        if len(search_results) >= 6:
                            break
                            
            except Exception as e:
                print(f"General web search attempt failed: {e}")
            
            return search_results
            
    except Exception as e:
        print(f"General web search URL retrieval failed: {e}")
        return []

def is_quality_source(url: str) -> bool:
    """Filter out low-quality or problematic sources."""
    low_quality_domains = [
        'pinterest.com', 'instagram.com', 'facebook.com', 'twitter.com',
        'tiktok.com', 'youtube.com', 'reddit.com'
    ]
    
    for domain in low_quality_domains:
        if domain in url.lower():
            return False
    
    return True

def has_sufficient_results(results: list) -> bool:
    """Check if search results are sufficient to proceed."""
    if not results:
        return False
    
    # Check for minimum number of results
    if len(results) < 2:
        return False
    
    # Check if results have meaningful content
    meaningful_results = 0
    for result in results:
        snippet = result.get('snippet', '')
        title = result.get('title', '')
        
        # Consider result meaningful if it has substantial content
        if len(snippet) > 50 or len(title) > 10:
            meaningful_results += 1
    
    return meaningful_results >= 2

def fetch_and_process_results(results: list, source_type: str) -> list:
    """Fetch and process webpage content from search results."""
    documents = []
    
    for result in results[:4]:  # Process top 4 results
        url = result.get('url', '')
        title = result.get('title', 'No title')
        
        print(f"📄 Fetching content from: {title}")
        content = fetch_webpage_content(url)
        
        if content and len(content.strip()) > 100:  # Ensure meaningful content
            doc = Document(
                page_content=content,
                metadata={
                    "source": url, 
                    "title": title,
                    "source_type": source_type
                }
            )
            documents.append(doc)
    
    return documents

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[:30000]
        
    except Exception as e:
        print(f"Failed to fetch content from {url}: {e}")
        return ""

def search_documents_with_vector_store(documents: list, query: str, source_type: str = "Web") -> 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 with source type indication
        results = []
        results.append(f"🔍 Search Results from {source_type}:\n")
        
        for i, doc in enumerate(relevant_docs, 1):
            source = doc.metadata.get('source', 'Unknown source')
            title = doc.metadata.get('title', 'No title')
            source_type_meta = doc.metadata.get('source_type', source_type)
            content = doc.page_content[:2000]  # Increased content length
            
            results.append(f"Result {i} ({source_type_meta}) - {title}:\n{content}\nSource: {source}\n")
        
        return "\n---\n".join(results)
        
    except Exception as e:
        return f"Vector search failed: {str(e)}"

web_search_tool = Tool(
    name="enhanced_web_retrieval",
    func=enhanced_web_retrieval_tool_func,
    description="Enhanced cascading web search with vector retrieval. First searches Wikipedia for reliable factual information, then falls back to general web search if insufficient results are found. Supports multiple search backends (auto, html, lite, bing) and uses semantic search to find relevant information. Ideal for comprehensive research on any topic."
)

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