# GAIA Dataset - Final Working Implementation for 300 Questions
# Based on actual GAIA dataset structure: 2023_all config with test split

import os
import json
import pandas as pd
import numpy as np
import time
from typing import Dict, List, Tuple
import requests
from datetime import datetime
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

# Install required packages
print("📦 Installing required packages...")
import subprocess
import sys

packages = [
    "langchain-community", "langchain-core", "langchain-google-genai", 
    "langchain-groq", "langchain-huggingface", "langgraph", "supabase", 
    "sentence-transformers", "tavily-python", "wikipedia", "arxiv", 
    "python-dotenv", "gradio", "datasets", "huggingface_hub"
]

for package in packages:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])

print("✅ All packages installed!")

# Import libraries
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_groq import ChatGroq
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import SupabaseVectorStore
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition, ToolNode
from supabase.client import Client, create_client
from datasets import load_dataset
from huggingface_hub import login, hf_hub_download

# ================================
# API KEYS SETUP
# ================================

print("🔑 Setting up API keys...")

API_KEYS = {
    'GROQ_API_KEY': "",
    'TAVILY_API_KEY': "",
    'SUPABASE_URL': '',
    'SUPABASE_SERVICE_KEY': '',
    'GOOGLE_API_KEY': '',
    'HUGGINGFACEHUB_API_TOKEN': 'h'
}

for key, value in API_KEYS.items():
    os.environ[key] = value

print(" API keys configured!")

# HuggingFace login
try:
    login(token=API_KEYS['HUGGINGFACEHUB_API_TOKEN'])
    print("HuggingFace authentication successful!")
except Exception as e:
    print(f"HuggingFace login warning: {e}")

# ================================
# FINAL GAIA DATASET LOADER
# ================================

class FinalGAIALoader:
    """Final GAIA dataset loader using correct configuration"""
    
    def __init__(self):
        self.dataset = None
        self.test_data = []
        
    def load_gaia_dataset(self):
        """Load GAIA dataset using the correct configuration"""
        print("Loading GAIA dataset with correct configuration...")
        
        try:
            # Based on the dataset code, use "2023_all" config
            print("Loading with config '2023_all'...")
            
            self.dataset = load_dataset(
                "gaia-benchmark/GAIA", 
                "2023_all",
                token=True,
                trust_remote_code=True  # Important for custom dataset scripts
            )
            
            print(f"Dataset loaded successfully!")
            print(f"Available splits: {list(self.dataset.keys())}")
            
            if 'test' in self.dataset:
                self.test_data = list(self.dataset['test'])
                print(f"Loaded {len(self.test_data)} test questions")
                
                # Analyze the data structure
                self._analyze_data_structure()
                return True
            else:
                print("No test split found")
                return False
                
        except Exception as e:
            print(f"Primary method failed: {e}")
            
            # Fallback: Try without trust_remote_code
            try:
                print("Trying fallback method...")
                self.dataset = load_dataset(
                    "gaia-benchmark/GAIA", 
                    "2023_all",
                    token=True
                )
                
                if 'test' in self.dataset:
                    self.test_data = list(self.dataset['test'])
                    print(f"Fallback successful! Loaded {len(self.test_data)} questions")
                    self._analyze_data_structure()
                    return True
                    
            except Exception as e2:
                print(f"Fallback also failed: {e2}")
                
                # Final fallback: Manual download
                return self._manual_download_fallback()
    
    def _manual_download_fallback(self):
        """Manual download fallback method"""
        print("Attempting manual download...")
        
        try:
            # Download the metadata file directly
            metadata_file = hf_hub_download(
                repo_id="gaia-benchmark/GAIA",
                filename="2023/test/metadata.jsonl",
                repo_type="dataset",
                token=API_KEYS['HUGGINGFACEHUB_API_TOKEN']
            )
            
            # Read the JSONL file
            test_questions = []
            with open(metadata_file, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.strip():
                        test_questions.append(json.loads(line))
            
            self.test_data = test_questions
            print(f"Manual download successful! Loaded {len(test_questions)} questions")
            self._analyze_data_structure()
            return True
            
        except Exception as e:
            print(f"Manual download failed: {e}")
            return False
    
    def _analyze_data_structure(self):
        """Analyze the structure of loaded GAIA data"""
        if not self.test_data:
            return
        
        print(f"\n GAIA Data Analysis:")
        print(f"Total questions: {len(self.test_data)}")
        
        # Show sample item structure
        if self.test_data:
            sample = self.test_data[0]
            print(f"Sample item keys: {list(sample.keys())}")
            
            # Level distribution
            levels = [item.get('Level') for item in self.test_data]
            level_counts = pd.Series(levels).value_counts().sort_index()
            print(f"\nLevel distribution:")
            for level, count in level_counts.items():
                print(f"  Level {level}: {count} questions")
            
            # Sample questions
            print(f"\nSample questions by level:")
            for level in sorted(level_counts.index):
                level_questions = [item for item in self.test_data if item.get('Level') == level]
                if level_questions:
                    sample_q = level_questions[0]['Question'][:100]
                    print(f"  Level {level}: {sample_q}...")
    
    def filter_300_questions(self, level1_ratio=0.6, level2_ratio=0.25, level3_ratio=0.15):
        """Filter exactly 300 questions with specified ratios"""
        
        if not self.test_data:
            print("No test data available")
            return []
        
        total_questions = 300
        level1_count = int(total_questions * level1_ratio)  # 180
        level2_count = int(total_questions * level2_ratio)  # 75
        level3_count = total_questions - level1_count - level2_count  # 45
        
        print(f"Filtering {total_questions} questions:")
        print(f"  Level 1: {level1_count} questions ({level1_ratio*100:.0f}%)")
        print(f"  Level 2: {level2_count} questions ({level2_ratio*100:.0f}%)")
        print(f"  Level 3: {level3_count} questions ({level3_ratio*100:.0f}%)")
        
        # Group by level
        level_groups = {1: [], 2: [], 3: []}
        for item in self.test_data:
            level = item.get('Level')
            if level in level_groups:
                level_groups[level].append(item)
        
        # Check availability
        print(f"\nAvailable questions by level:")
        for level in [1, 2, 3]:
            available = len(level_groups[level])
            print(f"  Level {level}: {available} available")
        
        # Sample questions
        filtered_questions = []
        np.random.seed(42)  # For reproducibility
        
        for level, target_count in [(1, level1_count), (2, level2_count), (3, level3_count)]:
            available = len(level_groups[level])
            
            if available >= target_count:
                # Random sample
                sampled_indices = np.random.choice(available, size=target_count, replace=False)
                sampled = [level_groups[level][i] for i in sampled_indices]
                filtered_questions.extend(sampled)
                print(f"Level {level}: Selected {target_count} from {available}")
            else:
                # Take all available
                filtered_questions.extend(level_groups[level])
                print(f"Level {level}: Only {available} available (needed {target_count})")
        
        print(f"\n📊 Total filtered: {len(filtered_questions)} questions")
        
        # Verify final distribution
        final_levels = [q['Level'] for q in filtered_questions]
        final_dist = pd.Series(final_levels).value_counts().sort_index()
        print(f"Final distribution:")
        for level, count in final_dist.items():
            percentage = (count / len(filtered_questions)) * 100
            print(f"  Level {level}: {count} questions ({percentage:.1f}%)")
        
        return filtered_questions
    
    def create_dataframe(self, questions):
        """Create DataFrame from filtered questions"""
        if not questions:
            return pd.DataFrame()
        
        data = []
        for i, item in enumerate(questions):
            data.append({
                'id': i,
                'task_id': item.get('task_id', f'gaia_{i}'),
                'question': item.get('Question', ''),
                'level': item.get('Level', 1),
                'final_answer': item.get('Final answer', ''),
                'file_name': item.get('file_name', ''),
                'annotator_metadata': item.get('Annotator Metadata', {})
            })
        
        df = pd.DataFrame(data)
        print(f"📋 Created DataFrame with {len(df)} questions")
        return df

# ================================
# ENHANCED SYSTEM PROMPT
# ================================

SYSTEM_PROMPT = """
You are an expert research assistant designed to answer complex, multi-step questions from the GAIA benchmark.
You have access to powerful tools for web search, Wikipedia, arXiv research, and mathematical calculations.

Key Instructions:
1. Read the question carefully and identify what information you need
2. Use tools strategically - web search for current info, Wikipedia for general knowledge, arXiv for research
3. For mathematical questions, use the calculation tools
4. Think step-by-step and show your reasoning
5. Be precise and accurate in your final answer

Answer Format:
- Provide clear reasoning and methodology
- Show any calculations or research steps
- End with: FINAL ANSWER: [exact answer]
- Keep the final answer concise and precise
- Match the format requested (number, name, list, etc.)

Examples:
Question: What is the population of Tokyo according to the latest census?
I need to search for the most recent population data for Tokyo.
[uses web_search("Tokyo population latest census")]
Based on the search results, the latest census shows...
FINAL ANSWER: 13960000

Question: What is 157 * 238?
I need to calculate 157 multiplied by 238.
[uses multiply(157, 238)]
FINAL ANSWER: 37366
"""

# ================================
# ENHANCED TOOLS
# ================================

@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 abs(b) < 1e-10:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def power(a: float, b: float) -> float:
    """Calculate a raised to the power of b."""
    return a ** b

@tool
def square_root(a: float) -> float:
    """Calculate the square root of a number."""
    if a < 0:
        raise ValueError("Cannot calculate square root of negative number.")
    return a ** 0.5

@tool
def modulo(a: float, b: float) -> float:
    """Calculate a modulo b (remainder after division)."""
    return a % b

@tool
def absolute_value(a: float) -> float:
    """Calculate the absolute value of a number."""
    return abs(a)

@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for comprehensive information on a topic."""
    try:
        search_docs = WikipediaLoader(query=query, load_max_docs=3).load()
        if not search_docs:
            return f"No Wikipedia results found for: {query}"
        
        # Return the most relevant result with more content
        result = search_docs[0]
        title = result.metadata.get('title', 'Unknown')
        content = result.page_content[:1500]
        
        return f"Wikipedia article '{title}':\n{content}..."
    except Exception as e:
        return f"Wikipedia search failed: {str(e)}"

@tool
def web_search(query: str) -> str:
    """Search the web for current, factual information."""
    try:
        search_tool = TavilySearchResults(max_results=4)
        results = search_tool.invoke(query)
        
        if not results:
            return f"No web results found for: {query}"
        
        formatted_results = f"Web search results for '{query}':\n\n"
        for i, result in enumerate(results, 1):
            url = result.get('url', 'Unknown URL')
            content = result.get('content', 'No content')[:800]
            formatted_results += f"{i}. Source: {url}\n{content}...\n\n"
        
        return formatted_results
    except Exception as e:
        return f"Web search failed: {str(e)}"

@tool
def arxiv_search(query: str) -> str:
    """Search arXiv for academic papers and research."""
    try:
        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
        if not search_docs:
            return f"No arXiv papers found for: {query}"
        
        formatted_results = f"arXiv search results for '{query}':\n\n"
        for i, doc in enumerate(search_docs, 1):
            title = doc.metadata.get('Title', 'Unknown Title')
            authors = doc.metadata.get('Authors', 'Unknown Authors')
            content = doc.page_content[:1000]
            formatted_results += f"{i}. Title: {title}\nAuthors: {authors}\nAbstract: {content}...\n\n"
        
        return formatted_results
    except Exception as e:
        return f"arXiv search failed: {str(e)}"

# All tools
tools = [
    multiply, add, subtract, divide, power, square_root, modulo, absolute_value,
    wiki_search, web_search, arxiv_search
]

# ================================
# ENHANCED GAIA AGENT
# ================================

class EnhancedGAIAAgent:
    """Enhanced GAIA agent optimized for benchmark performance"""
    
    def __init__(self, provider="groq", model="llama3-70b-8192"):
        self.provider = provider
        self.model = model
        self.graph = self._build_graph()
        
    def _build_graph(self):
        """Build optimized agent graph"""
        print(f"Building enhanced agent: {self.provider} {self.model}")
        
        if self.provider == "groq":
            llm = ChatGroq(
                model=self.model, 
                temperature=0,
                max_retries=3,
                timeout=120
            )
        elif self.provider == "google":
            llm = ChatGoogleGenerativeAI(
                model="gemini-pro", 
                temperature=0,
                max_retries=3
            )
        else:
            raise ValueError("Choose 'groq' or 'google'")
        
        llm_with_tools = llm.bind_tools(tools)
        
        def assistant_node(state: MessagesState):
            """Enhanced assistant node with better error handling"""
            try:
                messages = state["messages"]
                
                # Add system message if not present
                if not any(isinstance(msg, SystemMessage) for msg in messages):
                    messages = [SystemMessage(content=SYSTEM_PROMPT)] + messages
                
                # Invoke LLM
                response = llm_with_tools.invoke(messages)
                return {"messages": [response]}
                
            except Exception as e:
                error_msg = f"Assistant error: {str(e)}"
                print(f" {error_msg}")
                return {"messages": [HumanMessage(content=f"Error: {error_msg}")]}
        
        # Build graph
        builder = StateGraph(MessagesState)
        builder.add_node("assistant", assistant_node)
        builder.add_node("tools", ToolNode(tools))
        
        builder.add_edge(START, "assistant")
        builder.add_conditional_edges("assistant", tools_condition)
        builder.add_edge("tools", "assistant")
        
        return builder.compile()
    
    def process_question(self, question: str, question_id: str = None) -> Dict:
        """Process a single GAIA question with enhanced handling"""
        start_time = time.time()
        
        try:
            # Create message and invoke graph
            messages = [HumanMessage(content=question)]
            result = self.graph.invoke({"messages": messages})
            
            # Extract final response
            final_response = result["messages"][-1].content
            
            # Extract final answer more robustly
            answer = self._extract_final_answer(final_response)
            
            processing_time = time.time() - start_time
            
            return {
                'question_id': question_id,
                'question': question,
                'full_response': final_response,
                'final_answer': answer,
                'processing_time': processing_time,
                'status': 'success',
                'error': None
            }
            
        except Exception as e:
            processing_time = time.time() - start_time
            error_msg = str(e)
            
            return {
                'question_id': question_id,
                'question': question,
                'full_response': f"Error: {error_msg}",
                'final_answer': f"ERROR: {error_msg}",
                'processing_time': processing_time,
                'status': 'error',
                'error': error_msg
            }
    
    def _extract_final_answer(self, response: str) -> str:
        """Extract final answer more robustly"""
        if "FINAL ANSWER:" in response:
            answer = response.split("FINAL ANSWER:")[-1].strip()
            # Clean up the answer
            answer = answer.split('\n')[0].strip()  # Take first line only
            return answer
        else:
            # Fallback: take last substantial line
            lines = [line.strip() for line in response.split('\n') if line.strip()]
            return lines[-1] if lines else response.strip()
    
    def evaluate_300_questions(self, questions_df: pd.DataFrame) -> pd.DataFrame:
        """Evaluate 300 GAIA questions with comprehensive tracking"""
        print(f"Starting GAIA evaluation: {len(questions_df)} questions")
        print(f" Estimated time: {len(questions_df) * 15 / 60:.0f}-{len(questions_df) * 25 / 60:.0f} minutes")
        
        results = []
        
        # Progress tracking
        with tqdm(total=len(questions_df), desc="GAIA Evaluation Progress") as pbar:
            for idx, row in questions_df.iterrows():
                question_id = row.get('task_id', f'gaia_{idx}')
                question = row['question']
                level = row['level']
                expected = row.get('final_answer', '')
                
                print(f"\n Question {idx+1}/{len(questions_df)} - Level {level}")
                print(f"ID: {question_id}")
                print(f"Q: {question[:120]}...")
                
                # Process question
                result = self.process_question(question, question_id)
                
                # Add metadata
                result.update({
                    'level': level,
                    'expected_answer': expected,
                    'question_index': idx
                })
                
                results.append(result)
                
                # Show result
                status_emoji = "" if result['status'] == 'success' else ""
                print(f"A: {result['final_answer'][:120]}...")
                print(f"{status_emoji} Time: {result['processing_time']:.2f}s")
                
                # Update progress bar
                pbar.update(1)
                pbar.set_postfix({
                    'Success Rate': f"{len([r for r in results if r['status'] == 'success'])/len(results):.1%}",
                    'Avg Time': f"{np.mean([r['processing_time'] for r in results]):.1f}s"
                })
                
                # Save progress every 25 questions
                if (idx + 1) % 25 == 0:
                    temp_df = pd.DataFrame(results)
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    temp_df.to_csv(f'gaia_progress_{timestamp}.csv', index=False)
                    print(f"💾 Progress saved at question {idx+1}")
                
                # Rate limiting
                time.sleep(1.5)
        
        results_df = pd.DataFrame(results)
        print(f"\n GAIA evaluation completed!")
        return results_df

# ================================
# ANALYSIS FUNCTIONS
# ================================

def analyze_gaia_results(results_df: pd.DataFrame) -> Dict:
    """Comprehensive analysis of GAIA results"""
    
    if results_df.empty:
        return {}
    
    # Basic metrics
    total = len(results_df)
    successful = len(results_df[results_df['status'] == 'success'])
    success_rate = successful / total
    error_rate = 1 - success_rate
    
    metrics = {
        'total_questions': total,
        'successful_runs': successful,
        'success_rate': success_rate,
        'error_rate': error_rate,
        'avg_processing_time': results_df['processing_time'].mean(),
        'median_processing_time': results_df['processing_time'].median(),
        'total_processing_time': results_df['processing_time'].sum()
    }
    
    # Level-wise analysis
    level_metrics = {}
    for level in sorted(results_df['level'].unique()):
        level_data = results_df[results_df['level'] == level]
        level_success = len(level_data[level_data['status'] == 'success'])
        
        level_metrics[f'level_{level}'] = {
            'count': len(level_data),
            'success_count': level_success,
            'success_rate': level_success / len(level_data),
            'error_rate': 1 - (level_success / len(level_data)),
            'avg_time': level_data['processing_time'].mean()
        }
    
    metrics['by_level'] = level_metrics
    
    # Print comprehensive analysis
    print(f"\n GAIA Benchmark Results Analysis")
    print(f"=" * 60)
    print(f" Total Questions: {total}")
    print(f" Successful: {successful}")
    print(f" Overall Success Rate: {success_rate:.2%}")
    print(f" Error Rate: {error_rate:.2%}")
    print(f" Average Time: {metrics['avg_processing_time']:.2f}s")
    print(f" Total Time: {metrics['total_processing_time']/60:.1f} minutes")
    
    print(f"\n Performance by Difficulty Level:")
    for level_key, level_data in level_metrics.items():
        level_num = level_key.split('_')[1]
        print(f"Level {level_num}:")
        print(f"   Questions: {level_data['count']}")
        print(f"   Success: {level_data['success_count']}/{level_data['count']} ({level_data['success_rate']:.1%})")
        print(f"   Avg Time: {level_data['avg_time']:.1f}s")
    
    return metrics

def save_gaia_results(results_df: pd.DataFrame, metrics: Dict):
    """Save comprehensive GAIA results"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Save detailed results
    results_filename = f'gaia_300_final_results_{timestamp}.csv'
    results_df.to_csv(results_filename, index=False)
    print(f" Detailed results: {results_filename}")
    
    # Save metrics
    metrics_filename = f'gaia_300_metrics_{timestamp}.json'
    with open(metrics_filename, 'w') as f:
        json.dump(metrics, f, indent=2, default=str)
    print(f" Metrics: {metrics_filename}")
    
    # Create summary report
    report = f"""
# GAIA Benchmark - 300 Questions Evaluation Report
Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

## Overall Performance
- Total Questions: {metrics['total_questions']}
- Success Rate: {metrics['success_rate']:.2%}
- Average Time: {metrics['avg_processing_time']:.2f}s
- Total Time: {metrics['total_processing_time']/60:.1f} minutes

## Level Performance
"""
    
    for level_key, level_data in metrics['by_level'].items():
        level_num = level_key.split('_')[1]
        report += f"""
### Level {level_num}
- Questions: {level_data['count']}
- Success Rate: {level_data['success_rate']:.2%}
- Average Time: {level_data['avg_time']:.2f}s
"""
    
    report_filename = f'gaia_300_report_{timestamp}.txt'
    with open(report_filename, 'w') as f:
        f.write(report)
    print(f"💾 Report: {report_filename}")
    
    return timestamp

# ================================
# MAIN EXECUTION
# ================================

def main_gaia_300():
    """Main function for 300 GAIA questions evaluation"""
    print(" GAIA Benchmark - 300 Questions Evaluation")
    print("=" * 60)
    
    try:
        # Step 1: Load GAIA dataset
        print("\n Step 1: Loading GAIA Dataset")
        loader = FinalGAIALoader()
        
        if not loader.load_gaia_dataset():
            print(" Failed to load GAIA dataset")
            return None, None
        
        # Step 2: Filter 300 questions
        print(f"\n Step 2: Filtering 300 Questions")
        filtered_questions = loader.filter_300_questions()
        
        if len(filtered_questions) < 250:  # Allow some flexibility
            print(f" Only {len(filtered_questions)} questions available")
            proceed = input("Proceed with available questions? (y/n): ")
            if proceed.lower() != 'y':
                return None, None
        
        # Create DataFrame
        questions_df = loader.create_dataframe(filtered_questions)
        
        # Step 3: Confirm evaluation
        print(f"\n Step 3: Ready for Evaluation")
        print(f" Questions to evaluate: {len(questions_df)}")
        
        level_dist = questions_df['level'].value_counts().sort_index()
        for level, count in level_dist.items():
            percentage = (count / len(questions_df)) * 100
            print(f"  Level {level}: {count} questions ({percentage:.1f}%)")
        
        estimated_time = len(questions_df) * 18 / 60  # 18 seconds average per question
        print(f" Estimated time: {estimated_time:.0f} minutes")
        
        # Confirm with user
        proceed = input(f"\nProceed with {len(questions_df)} questions evaluation? (y/n): ")
        if proceed.lower() != 'y':
            print("Evaluation cancelled.")
            return None, None
        
        # Step 4: Initialize enhanced agent
        print(f"\n Step 4: Initializing Enhanced GAIA Agent")
        agent = EnhancedGAIAAgent(provider="groq", model="llama3-70b-8192")
        
        # Step 5: Run evaluation
        print(f"\n Step 5: Running GAIA Evaluation")
        results_df = agent.evaluate_300_questions(questions_df)
        
        # Step 6: Analyze results
        print(f"\n Step 6: Analyzing Results")
        metrics = analyze_gaia_results(results_df)
        
        # Step 7: Save results
        print(f"\n Step 7: Saving Results")
        timestamp = save_gaia_results(results_df, metrics)
        
        # Final summary
        print(f"\n GAIA Evaluation Completed Successfully!")
        print(f" Success Rate: {metrics['success_rate']:.2%}")
        print(f" Total Time: {metrics['total_processing_time']/60:.1f} minutes")
        print(f" Files saved with timestamp: {timestamp}")
        
        # Performance insights
        best_level = max(metrics['by_level'].items(), key=lambda x: x[1]['success_rate'])
        worst_level = min(metrics['by_level'].items(), key=lambda x: x[1]['success_rate'])
        
        print(f"\n Key Insights:")
        print(f" Best Performance: Level {best_level[0].split('_')[1]} ({best_level[1]['success_rate']:.1%})")
        print(f" Most Challenge: Level {worst_level[0].split('_')[1]} ({worst_level[1]['success_rate']:.1%})")
        
        if metrics['success_rate'] >= 0.85:
            print(" Excellent performance on GAIA benchmark!")
        elif metrics['success_rate'] >= 0.70:
            print(" Good performance on GAIA benchmark!")
        else:
            print(" Room for improvement on GAIA benchmark.")
        
        return results_df, metrics
        
    except Exception as e:
        print(f" Error in main execution: {e}")
        import traceback
        traceback.print_exc()
        return None, None

def test_gaia_single():
    """Test with a single GAIA question"""
    print(" Testing Single GAIA Question")
    print("=" * 40)
    
    try:
        # Load dataset
        loader = FinalGAIALoader()
        if not loader.load_gaia_dataset():
            print(" Failed to load dataset")
            return None
        
        # Get one question of each level
        sample_questions = []
        for level in [1, 2, 3]:
            level_questions = [q for q in loader.test_data if q.get('Level') == level]
            if level_questions:
                sample_questions.append(level_questions[0])
        
        if not sample_questions:
            print(" No sample questions found")
            return None
        
        # Test with agent
        agent = EnhancedGAIAAgent()
        
        for i, question_data in enumerate(sample_questions):
            question = question_data['Question']
            level = question_data['Level']
            
            print(f"\n Test {i+1} - Level {level}")
            print(f"Q: {question[:150]}...")
            
            result = agent.process_question(question, f"test_{level}")
            
            print(f"A: {result['final_answer']}")
            print(f" Time: {result['processing_time']:.2f}s")
            print(f"Status: {result['status']}")
        
        return True
        
    except Exception as e:
        print(f" Test failed: {e}")
        return None

def quick_gaia_test():
    """Quick test with 10 questions"""
    print(" Quick GAIA Test - 10 Questions")
    print("=" * 40)
    
    try:
        # Load and filter
        loader = FinalGAIALoader()
        if not loader.load_gaia_dataset():
            return None, None
        
        # Get 10 questions (3/4/3 distribution)
        level_groups = {1: [], 2: [], 3: []}
        for item in loader.test_data:
            level = item.get('Level')
            if level in level_groups:
                level_groups[level].append(item)
        
        quick_questions = []
        quick_questions.extend(level_groups[1][:3])  # 3 Level 1
        quick_questions.extend(level_groups[2][:4])  # 4 Level 2  
        quick_questions.extend(level_groups[3][:3])  # 3 Level 3
        
        questions_df = loader.create_dataframe(quick_questions)
        
        # Run evaluation
        agent = EnhancedGAIAAgent()
        results_df = agent.evaluate_300_questions(questions_df)
        
        # Quick analysis
        metrics = analyze_gaia_results(results_df)
        
        return results_df, metrics
        
    except Exception as e:
        print(f" Quick test failed: {e}")
        return None, None

# ================================
# EXECUTION
# ================================

if __name__ == "__main__":
    print(" GAIA Benchmark Evaluation System")
    print("Using correct dataset configuration: 2023_all")
    print("=" * 60)
    
    # Test API connectivity
    print("\n Testing API Connectivity...")
    try:
        # Test Groq
        test_llm = ChatGroq(model="llama3-8b-8192", temperature=0)
        test_response = test_llm.invoke([HumanMessage(content="Hello")])
        print(" Groq API working")
        
        # Test Tavily
        test_search = TavilySearchResults(max_results=1)
        test_results = test_search.invoke("test")
        print(" Tavily API working")
        
    except Exception as e:
        print(f" API test warning: {e}")
    
    # Main menu
    print(f"\n🎮 Choose evaluation option:")
    print("1. Test single GAIA questions (5 minutes)")
    print("2. Quick test with 10 questions (15 minutes)")
    print("3. Full 300 questions evaluation (75-90 minutes)")
    print("4. Auto-run full evaluation")
    
    choice = input("\nEnter choice (1-4): ").strip()
    
    if choice == "1":
        print(" Running single question tests...")
        test_gaia_single()
        
    elif choice == "2":
        print(" Running quick 10-question test...")
        results_df, metrics = quick_gaia_test()
        if results_df is not None:
            print(f"\n Quick Test Results:")
            print(f" Success Rate: {metrics['success_rate']:.2%}")
            print(f"⏱️ Average Time: {metrics['avg_processing_time']:.1f}s")
        
    elif choice == "3" or choice == "4":
        print(" Starting full GAIA 300 questions evaluation...")
        results_df, metrics = main_gaia_300()
        
        if results_df is not None:
            print(f"\n Final GAIA Benchmark Results:")
            print(f" Overall Success Rate: {metrics['success_rate']:.2%}")
            print(f" Questions Completed: {len(results_df)}")
            print(f"⏱️ Total Evaluation Time: {metrics['total_processing_time']/60:.1f} minutes")
            
            # Level breakdown
            for level_key, level_data in metrics['by_level'].items():
                level_num = level_key.split('_')[1]
                print(f" Level {level_num}: {level_data['success_rate']:.1%} success rate")
                
        else:
            print(" Evaluation failed")
            
    else:
        print("Invalid choice. Running single question test...")
        test_gaia_single()

print("\n🏁 GAIA Evaluation System Complete!")