from enum import Enum
from typing import Literal

from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmarkConfig


class VisualWebBenchTask(str, Enum):
    """Task for VisualWebBench"""

    CAPTION_TASK = "web_caption"
    HEADING_OCR_TASK = "heading_ocr"
    WEBQA_TASK = "webqa"
    ELEMENT_OCR_TASK = "element_ocr"
    ACTION_PREDICTION_TASK = "action_prediction"
    ELEMENT_GROUND_TASK = "element_ground"
    ACTION_GROUND_TASK = "action_ground"


class VisualWebBenchConfig(HubBaseBenchmarkConfig):
    hf_repo: Literal["visualwebbench/VisualWebBench"] = "visualwebbench/VisualWebBench"
    """HF repo name"""

    revision: str = "main"

    split: Literal["test"] = "test"
    """HF split name"""

    max_tokens: int = 512
    """Maximum number of tokens in the completion."""

    temperature: float = 0.0
    """Sampling temperature."""

    data_dir: str | list[str] | None = [task.value for task in VisualWebBenchTask]
    """HF data directory"""

    @classmethod
    def create(cls) -> "VisualWebBenchConfig":
        """Create a config for ScreenSpot v1 training dataset"""
        return cls()
