from fastapi import APIRouter, HTTPException
import random
from datasets import load_dataset
from huggingface_hub import HfApi, dataset_info
import os

router = APIRouter(tags=["benchmark"])

@router.get("/benchmark-questions/{session_id}")
async def get_benchmark_questions(session_id: str):
    """
    Get example questions from the generated benchmark
    
    Args:
        session_id: Session ID for the benchmark
        
    Returns:
        Dictionary with sample questions from the dataset
    """
    try:
        # Dataset path on Hugging Face
        dataset_repo_id = f"yourbench/yourbench_{session_id}"
        
        # Initialize response
        response = {
            "success": False,
            "questions": [],
            "dataset_url": f"https://huggingface.co/datasets/{dataset_repo_id}"
        }
        
        # Try to load the dataset
        questions = []
        
        try:
            # Essayer de charger les questions single-shot directement avec le nom de config
            single_dataset = load_dataset(dataset_repo_id, 'single_shot_questions')
            if single_dataset and len(single_dataset['train']) > 0:
                # Get a random sample (up to 2) from single-shot questions
                sample_indices = random.sample(range(len(single_dataset['train'])), min(2, len(single_dataset['train'])))
                for idx in sample_indices:
                    questions.append({
                        "id": str(idx),
                        "question": single_dataset['train'][idx].get("question", ""),
                        "answer": single_dataset['train'][idx].get("self_answer", "No answer available"),
                        "type": "single_shot"
                    })
                print(f"Loaded {len(questions)} single-shot questions")
        except Exception as e:
            print(f"Error loading single-shot questions: {str(e)}")
        
        try:
            # Essayer de charger les questions multi-hop si nécessaire
            if len(questions) < 2:
                multi_dataset = load_dataset(dataset_repo_id, 'multi_hop_questions')
                if multi_dataset and len(multi_dataset['train']) > 0:
                    # Get remaining questions from multi-hop questions
                    remaining = 2 - len(questions)
                    sample_indices = random.sample(range(len(multi_dataset['train'])), min(remaining, len(multi_dataset['train'])))
                    for idx in sample_indices:
                        questions.append({
                            "id": str(idx),
                            "question": multi_dataset['train'][idx].get("question", ""),
                            "answer": multi_dataset['train'][idx].get("self_answer", "No answer available"),
                            "type": "multi_hop"
                        })
                    print(f"Loaded {len(questions)} multi-hop questions")
        except Exception as e:
            print(f"Error loading multi-hop questions: {str(e)}")
        
        # If we couldn't load any questions, the dataset might not exist
        if len(questions) == 0:
            # Check if we have a directory for this session locally as fallback
            session_dir = os.path.join("uploaded_files", session_id)
            if not os.path.exists(session_dir):
                raise HTTPException(status_code=404, detail="Dataset not found")
        
        # Update the response
        response["success"] = len(questions) > 0
        response["questions"] = questions
        return response
            
    except HTTPException:
        # Re-raise HTTP exceptions
        raise
    except Exception as e:
        return {
            "success": False,
            "error": f"Error retrieving benchmark questions: {str(e)}"
        }