from enum import Enum
from typing import Literal

from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmarkConfig


class AndroidControlTask(str, Enum):
    """Task for AndroidControl"""

    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 AndroidControlConfig(HubBaseBenchmarkConfig):
    hf_repo: Literal["smolagents/android-control"] = "smolagents/android-control"
    """HF repo name"""

    revision: str = "main"

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

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

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

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

    max_step_memory: int = 10
    """Maximum number of steps."""

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