"""Tutorials - Evaluation Datasets - Example 2 - Loading from HuggingFace."""

import asyncio
from pathlib import Path
from pprint import pprint
from typing import Any, List

from dotenv import load_dotenv
from openai import AsyncOpenAI

from tutorials.utils import save_results_to_json, setup_logging

from scorebook import EvalDataset, evaluate_async


async def main() -> Any:
    """Run evaluations using datasets loaded from HuggingFace Hub.

    This example demonstrates how to load evaluation datasets directly from
    HuggingFace Hub using the from_huggingface method. This allows you to
    easily evaluate on standard benchmarks and datasets.

    We'll evaluate on the SimpleQA dataset, which tests factual question answering.

    Prerequisites:
        - OpenAI API key set in environment variable OPENAI_API_KEY
    """

    # Initialize OpenAI client
    client = AsyncOpenAI()
    model_name = "gpt-4o-mini"

    # Define an async inference function
    async def inference(inputs: List[Any], **hyperparameters: Any) -> List[Any]:
        """Process inputs through OpenAI's API.

        Args:
            inputs: Input values from an EvalDataset.
            hyperparameters: Model hyperparameters.

        Returns:
            List of model outputs for all inputs.
        """
        outputs = []
        for input_val in inputs:
            # Build messages for OpenAI API
            messages = [
                {
                    "role": "system",
                    "content": "Answer the question directly and concisely. Provide only the answer, no additional context or text.",
                },
                {"role": "user", "content": str(input_val)},
            ]

            # Call OpenAI API
            try:
                response = await client.chat.completions.create(
                    model=model_name,
                    messages=messages,
                    temperature=0.7,
                )
                output = response.choices[0].message.content.strip()
            except Exception as e:
                output = f"Error: {str(e)}"

            outputs.append(output)

        return outputs

    # Load dataset from HuggingFace Hub
    simple_qa = EvalDataset.from_huggingface(
        path="basicv8vc/SimpleQA",
        metrics="accuracy",
        input="problem",
        label="answer",
        split="test",
    )
    print(f"Loaded {simple_qa.name} from HuggingFace Hub: {len(simple_qa.items)} items")

    # Run evaluation with a sample to avoid long runtime
    results = await evaluate_async(
        inference,
        simple_qa,
        sample_size=10,  # Sample 10 items for quick testing
        return_aggregates=True,
        return_items=True,
        return_output=True,
        upload_results=False,
    )

    pprint(results)
    return results


if __name__ == "__main__":
    load_dotenv()
    log_file = setup_logging(experiment_id="2-evaluation_datasets_from_huggingface", base_dir=Path(__file__).parent)
    output_dir = Path(__file__).parent / "results"
    output_dir.mkdir(exist_ok=True)
    results_dict = asyncio.run(main())
    save_results_to_json(results_dict, output_dir, "2-evaluation_datasets_from_huggingface_output.json")
