from datetime import datetime
from typing import Literal

from inspect_ai import Task, task
from inspect_ai.model import GenerateConfig
from inspect_ai.solver import generate, system_message

from inspect_evals.livebench.scorer import decide_scorer
from inspect_evals.livebench.utils import aggregate_datasets, ensure_nltk_resource

LIVEBENCH_DATASETS = [
    "math",
    "reasoning",
    "coding",
    "language",
    "data_analysis",
    "instruction_following",
]

# Taken from Fastchat as per paper
# https://github.com/lm-sys/FastChat/blob/main/fastchat/llm_judge/data/judge_prompts.jsonl
SYSTEM_PROMPT = "You are a helpful assistant."


@task
def livebench(
    livebench_release_date: datetime | None = None,
    category: Literal[
        "math",
        "reasoning",
        "coding",
        "language",
        "data_analysis",
        "instruction_following",
    ]
    | None = None,
) -> Task:
    # Can be downloaded manually here: https://www.nltk.org/nltk_data/
    # Used for instruction following scoring on the livebench files below
    # https://github.com/LiveBench/LiveBench/blob/main/livebench/if_runner/instruction_following_eval/instructions_util.py
    # https://github.com/LiveBench/LiveBench/blob/main/livebench/if_runner/instruction_following_eval/instructions.py

    ensure_nltk_resource()

    ds = aggregate_datasets(
        LIVEBENCH_DATASETS,
        livebench_release_date=livebench_release_date,
        category=category,
    )

    return Task(
        dataset=ds,
        solver=[system_message(SYSTEM_PROMPT), generate()],
        scorer=[decide_scorer()],
        config=GenerateConfig(temperature=0),
    )
