"""LangGraph Agent - Enhanced for GAIA Benchmark"""
import os
import json
import re
import base64
from io import BytesIO
from typing import Dict, Any, List
from dotenv import load_dotenv
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
from langchain_community.vectorstores import SupabaseVectorStore
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool
from supabase.client import Client, create_client
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import requests
from urllib.parse import urlparse, parse_qs
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

load_dotenv()

# CSV Knowledge Base Search Tool - HIGHEST PRIORITY
@tool
def search_csv_knowledge_base(query: str, top_k: int = 3) -> str:
    """Search the CSV knowledge base for relevant questions and answers.
    This should be used FIRST for any question to check existing knowledge.
    
    Args:
        query: The search query/question
        top_k: Number of top results to return (default: 3)
    """
    try:
        # Load the CSV file
        df = pd.read_csv('supabase_docs.csv')
        
        if df.empty:
            return "CSV knowledge base is empty."
        
        # Extract questions and answers from content
        questions = []
        answers = []
        
        for content in df['content']:
            if pd.isna(content):
                continue
                
            content_str = str(content)
            
            # Extract question and answer using regex
            question_match = re.search(r'Question\s*:\s*(.*?)(?=Final answer|$)', content_str, re.IGNORECASE | re.DOTALL)
            answer_match = re.search(r'Final answer\s*:\s*(.*?)$', content_str, re.IGNORECASE | re.DOTALL)
            
            if question_match and answer_match:
                question = question_match.group(1).strip()
                answer = answer_match.group(1).strip()
                questions.append(question)
                answers.append(answer)
        
        if not questions:
            return "No valid questions found in CSV knowledge base."
        
        # Create TF-IDF vectorizer and compute similarities
        vectorizer = TfidfVectorizer(stop_words='english', ngram_range=(1, 2))
        
        # Combine query with questions for vectorization
        all_texts = [query] + questions
        tfidf_matrix = vectorizer.fit_transform(all_texts)
        
        # Calculate cosine similarities between query and questions
        query_vector = tfidf_matrix[0:1]
        question_vectors = tfidf_matrix[1:]
        similarities = cosine_similarity(query_vector, question_vectors).flatten()
        
        # Get top-k most similar questions
        top_indices = similarities.argsort()[-top_k:][::-1]
        
        results = []
        for i, idx in enumerate(top_indices):
            similarity_score = similarities[idx]
            
            # Only include results with reasonable similarity (threshold: 0.1)
            if similarity_score > 0.1:
                results.append({
                    'rank': i + 1,
                    'question': questions[idx],
                    'answer': answers[idx],
                    'similarity': similarity_score
                })
        
        if not results:
            return "No relevant results found in CSV knowledge base."
        
        # Format results
        formatted_results = []
        for result in results:
            formatted_results.append(
                f"**Result {result['rank']} (Similarity: {result['similarity']:.3f})**\n"
                f"Question: {result['question']}\n"
                f"Answer: {result['answer']}\n"
                f"{'='*50}"
            )
        
        return f"Found {len(results)} relevant results from CSV knowledge base:\n\n" + "\n\n".join(formatted_results)
        
    except FileNotFoundError:
        return "CSV knowledge base file 'supabase_docs.csv' not found."
    except Exception as e:
        return f"Error searching CSV knowledge base: {str(e)}"

# Enhanced Mathematical Tools
@tool
def calculate(expression: str) -> float:
    """Safely evaluate mathematical expressions including complex calculations.
    
    Args:
        expression: Mathematical expression as string (e.g., "2 + 3 * 4", "sqrt(16)", "log(100)")
    """
    try:
        # Safe evaluation with math functions
        import math
        safe_dict = {
            "__builtins__": {},
            "abs": abs, "round": round, "min": min, "max": max,
            "sum": sum, "pow": pow, "sqrt": math.sqrt, "log": math.log,
            "log10": math.log10, "exp": math.exp, "sin": math.sin,
            "cos": math.cos, "tan": math.tan, "pi": math.pi, "e": math.e
        }
        result = eval(expression, safe_dict)
        return float(result)
    except Exception as e:
        return f"Error calculating {expression}: {str(e)}"

@tool
def percentage_calculation(part: float, whole: float, operation: str = "find_percentage") -> float:
    """Calculate percentages and percentage-related operations.
    
    Args:
        part: The part value
        whole: The whole value  
        operation: Type of calculation - "find_percentage", "find_part", "find_whole"
    """
    try:
        if operation == "find_percentage":
            return (part / whole) * 100
        elif operation == "find_part":
            return (part / 100) * whole
        elif operation == "find_whole":
            return part / (whole / 100)
        else:
            return "Invalid operation. Use: find_percentage, find_part, or find_whole"
    except ZeroDivisionError:
        return "Error: Division by zero"
    except Exception as e:
        return f"Error: {str(e)}"

# Enhanced Search Tools
@tool
def advanced_web_search(query: str, num_results: int = 5) -> str:
    """Enhanced web search with more comprehensive results.
    
    Args:
        query: The search query
        num_results: Number of results to return (1-10)
    """
    try:
        search_tool = TavilySearchResults(max_results=min(num_results, 10))
        results = search_tool.invoke(query)
        
        formatted_results = []
        for i, result in enumerate(results, 1):
            content = result.get('content', '')
            url = result.get('url', '')
            title = result.get('title', f'Result {i}')
            
            formatted_results.append(
                f"**Result {i}: {title}**\n"
                f"URL: {url}\n"
                f"Content: {content}\n"
                f"{'='*50}"
            )
        
        return "\n\n".join(formatted_results)
    except Exception as e:
        return f"Search error: {str(e)}"

@tool
def wikipedia_detailed_search(query: str, sections: bool = True) -> str:
    """Enhanced Wikipedia search with section information.
    
    Args:
        query: Search query
        sections: Whether to include section information
    """
    try:
        docs = WikipediaLoader(query=query, load_max_docs=3).load()
        
        formatted_docs = []
        for doc in docs:
            content = doc.page_content
            source = doc.metadata.get("source", "Unknown")
            title = doc.metadata.get("title", "Unknown")
            
            # Truncate very long content
            if len(content) > 2000:
                content = content[:2000] + "... [truncated]"
            
            formatted_docs.append(
                f"**Wikipedia Article: {title}**\n"
                f"Source: {source}\n"
                f"Content: {content}\n"
                f"{'='*50}"
            )
        
        return "\n\n".join(formatted_docs)
    except Exception as e:
        return f"Wikipedia search error: {str(e)}"

# Date and Time Tools
@tool
def date_calculator(date_str: str, operation: str, days: int = 0, months: int = 0, years: int = 0) -> str:
    """Calculate dates and time differences.
    
    Args:
        date_str: Date in format YYYY-MM-DD or common formats
        operation: "add", "subtract", "days_between", "format"
        days: Number of days to add/subtract
        months: Number of months to add/subtract  
        years: Number of years to add/subtract
    """
    try:
        from dateutil.parser import parse
        from dateutil.relativedelta import relativedelta
        
        base_date = parse(date_str)
        
        if operation == "add":
            result_date = base_date + relativedelta(years=years, months=months, days=days)
            return result_date.strftime("%Y-%m-%d")
        elif operation == "subtract":
            result_date = base_date - relativedelta(years=years, months=months, days=days)
            return result_date.strftime("%Y-%m-%d")
        elif operation == "format":
            return base_date.strftime("%Y-%m-%d")
        else:
            return f"Unknown operation: {operation}"
    except Exception as e:
        return f"Date calculation error: {str(e)}"

@tool
def current_date_time() -> str:
    """Get current date and time information."""
    now = datetime.now()
    return {
        "current_date": now.strftime("%Y-%m-%d"),
        "current_time": now.strftime("%H:%M:%S"),
        "current_datetime": now.strftime("%Y-%m-%d %H:%M:%S"),
        "day_of_week": now.strftime("%A"),
        "month": now.strftime("%B"),
        "year": now.year
    }

# Data Analysis Tools
@tool
def analyze_numbers(numbers: List[float], operation: str = "all") -> Dict[str, float]:
    """Analyze a list of numbers with statistical operations.
    
    Args:
        numbers: List of numbers to analyze
        operation: "mean", "median", "mode", "std", "min", "max", "sum", "all"
    """
    try:
        import statistics
        
        if not numbers:
            return {"error": "Empty list provided"}
        
        results = {}
        
        if operation in ["mean", "all"]:
            results["mean"] = statistics.mean(numbers)
        if operation in ["median", "all"]:
            results["median"] = statistics.median(numbers)
        if operation in ["mode", "all"]:
            try:
                results["mode"] = statistics.mode(numbers)
            except statistics.StatisticsError:
                results["mode"] = "No unique mode"
        if operation in ["std", "all"]:
            if len(numbers) > 1:
                results["standard_deviation"] = statistics.stdev(numbers)
        if operation in ["min", "all"]:
            results["minimum"] = min(numbers)
        if operation in ["max", "all"]:
            results["maximum"] = max(numbers)
        if operation in ["sum", "all"]:
            results["sum"] = sum(numbers)
        if operation in ["count", "all"]:
            results["count"] = len(numbers)
        
        return results
    except Exception as e:
        return {"error": f"Analysis error: {str(e)}"}

# Text Processing Tools
@tool
def extract_numbers_from_text(text: str) -> List[float]:
    """Extract all numbers from text.
    
    Args:
        text: Text to extract numbers from
    """
    try:
        # Find all numbers (including decimals and negatives)
        number_pattern = r'-?\d+\.?\d*'
        matches = re.findall(number_pattern, text)
        numbers = [float(match) for match in matches if match]
        return numbers
    except Exception as e:
        return [f"Error extracting numbers: {str(e)}"]

@tool
def text_analysis(text: str, analysis_type: str = "basic") -> Dict[str, Any]:
    """Analyze text for various properties.
    
    Args:
        text: Text to analyze
        analysis_type: "basic", "words", "sentences", "characters"
    """
    try:
        results = {}
        
        if analysis_type in ["basic", "characters"]:
            results["character_count"] = len(text)
            results["character_count_no_spaces"] = len(text.replace(" ", ""))
        
        if analysis_type in ["basic", "words"]:
            words = text.split()
            results["word_count"] = len(words)
            results["unique_words"] = len(set(word.lower() for word in words))
        
        if analysis_type in ["basic", "sentences"]:
            sentences = text.split('.')
            results["sentence_count"] = len([s for s in sentences if s.strip()])
        
        return results
    except Exception as e:
        return {"error": f"Text analysis error: {str(e)}"}

# URL and Web Tools
@tool
def analyze_url(url: str) -> Dict[str, str]:
    """Analyze and extract information from URLs.
    
    Args:
        url: URL to analyze
    """
    try:
        parsed = urlparse(url)
        query_params = parse_qs(parsed.query)
        
        return {
            "scheme": parsed.scheme,
            "domain": parsed.netloc,
            "path": parsed.path,
            "query_parameters": dict(query_params),
            "fragment": parsed.fragment
        }
    except Exception as e:
        return {"error": f"URL analysis error: {str(e)}"}

# Original tools (keeping the working ones)
@tool
def multiply(a: float, b: float) -> float:
    """Multiply two numbers."""
    return a * b

@tool
def add(a: float, b: float) -> float:
    """Add two numbers."""
    return a + b

@tool
def subtract(a: float, b: float) -> float:
    """Subtract two numbers."""
    return a - b

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

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

@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 results."""
    try:
        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 formatted_search_docs
    except Exception as e:
        return f"Wikipedia search error: {str(e)}"

@tool
def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 3 results."""
    try:
        search_docs = TavilySearchResults(max_results=3).invoke(query=query)
        formatted_search_docs = "\n\n---\n\n".join(
            [
                f'<Document source="{doc.get("url", "Unknown")}">\n{doc.get("content", "No content")}\n</Document>'
                for doc in search_docs
            ])
        return formatted_search_docs
    except Exception as e:
        return f"Web search error: {str(e)}"

@tool
def arxiv_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 results."""
    try:
        search_docs = ArxivLoader(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{doc.page_content[:1000]}\n</Document>'
                for doc in search_docs
            ])
        return formatted_search_docs
    except Exception as e:
        return f"Arxiv search error: {str(e)}"

# Load the enhanced system prompt
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()

# System message
sys_msg = SystemMessage(content=system_prompt)

# Build a retriever
try:
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
    supabase: Client = create_client(
        os.environ.get("SUPABASE_URL"), 
        os.environ.get("SUPABASE_SERVICE_KEY"))
    vector_store = SupabaseVectorStore(
        client=supabase,
        embedding=embeddings,
        table_name="documents",
        query_name="match_documents_langchain",
    )
    retriever_tool = create_retriever_tool(
        retriever=vector_store.as_retriever(),
        name="similar_questions_search",
        description="Search for similar questions and answers from a knowledge base to help with current question.",
    )
    use_retriever = True
except Exception as e:
    print(f"Warning: Could not initialize retriever: {e}")
    retriever_tool = None
    use_retriever = False

# Enhanced tools list
tools = [
    # CSV Knowledge Base Search Tool
    search_csv_knowledge_base,
    
    # Mathematical tools
    calculate,
    percentage_calculation,
    multiply,
    add,
    subtract,
    divide,
    modulus,
    
    # Search tools
    advanced_web_search,
    wikipedia_detailed_search,
    web_search,
    wiki_search,
    arxiv_search,
    
    # Date and time tools
    date_calculator,
    current_date_time,
    
    # Data analysis tools
    analyze_numbers,
    extract_numbers_from_text,
    text_analysis,
    
    # Web tools
    analyze_url,
]

# Add retriever tool if available
if use_retriever and retriever_tool:
    tools.append(retriever_tool)

def build_graph(provider: str = "groq"):
    """Build the enhanced graph for GAIA benchmark."""
    
    if provider == "google":
        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
    elif provider == "groq":
        llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
    elif provider == "huggingface":
        llm = ChatHuggingFace(
            llm=HuggingFaceEndpoint(
                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
                temperature=0,
            ),
        )
    else:
        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
    
    # Bind tools to LLM
    llm_with_tools = llm.bind_tools(tools)

    def assistant(state: MessagesState):
        """Enhanced assistant node with better reasoning."""
        messages = state["messages"]
        
        # Add reasoning prompt to encourage step-by-step thinking
        if len(messages) > 0 and not any("Let me think through this step by step" in str(msg.content) for msg in messages):
            reasoning_prompt = HumanMessage(
                content="Let me think through this step by step and use the available tools to find the accurate answer."
            )
            messages = messages + [reasoning_prompt]
        
        response = llm_with_tools.invoke(messages)
        return {"messages": [response]}
    
    def retriever(state: MessagesState):
        """Enhanced retriever node."""
        messages = state["messages"]
        if use_retriever and len(messages) > 0:
            try:
                query = messages[0].content
                similar_questions = vector_store.similarity_search(query, k=2)
                
                if similar_questions:
                    examples = "\n\n".join([
                        f"Example {i+1}:\n{doc.page_content}" 
                        for i, doc in enumerate(similar_questions)
                    ])
                    
                    example_msg = HumanMessage(
                        content=f"Here are similar questions and answers for reference:\n\n{examples}"
                    )
                    return {"messages": [sys_msg] + messages + [example_msg]}
            except Exception as e:
                print(f"Retriever error: {e}")
        
        return {"messages": [sys_msg] + messages}

    # Build the graph
    builder = StateGraph(MessagesState)
    builder.add_node("retriever", retriever)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    
    builder.add_edge(START, "retriever")
    builder.add_edge("retriever", "assistant")
    builder.add_conditional_edges(
        "assistant",
        tools_condition,
    )
    builder.add_edge("tools", "assistant")

    return builder.compile()

# Test function
if __name__ == "__main__":
    question = "What is 15% of 240?"
    graph = build_graph(provider="groq")
    messages = [HumanMessage(content=question)]
    result = graph.invoke({"messages": messages})
    for m in result["messages"]:
        m.pretty_print()