from enum import Enum
from typing import Literal

from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmarkConfig
from screensuite.benchmarks.perception.screenspot.prompt import LocalizationPrompt


class ScreenSpotDatasetVersion(str, Enum):
    """Dataset version"""

    V1 = "v1"
    V2 = "v2"
    PRO = "pro"


class ScreenSpotHFRepo(str, Enum):
    """HF repo name"""

    V1 = "rootsautomation/ScreenSpot"
    V2 = "HongxinLi/ScreenSpot_v2"
    PRO = "HongxinLi/ScreenSpot-Pro"


class ScreenSpotConfig(HubBaseBenchmarkConfig):
    dataset_version: str
    """Dataset version"""

    hf_repo: str
    """HF repo name"""

    revision: str = "main"

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

    # Completion parameters
    system_prompt: LocalizationPrompt

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

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

    accuracy_threshold: float = 0.5
    """Accuracy threshold for correct predictions for bounding box prompt (IoU >= accuracy_threshold). For click prompt is always 1.0 or 0.0."""

    @classmethod
    def v1(cls, system_prompt: LocalizationPrompt) -> "ScreenSpotConfig":
        """Create a config for ScreenSpot v1 training dataset"""
        return cls(
            dataset_version=ScreenSpotDatasetVersion.V1.value,
            hf_repo=ScreenSpotHFRepo.V1.value,
            system_prompt=system_prompt,
        )

    @classmethod
    def v2(cls, system_prompt: LocalizationPrompt) -> "ScreenSpotConfig":
        """Create a config for ScreenSpot v2 training dataset"""
        return cls(
            dataset_version=ScreenSpotDatasetVersion.V2.value,
            hf_repo=ScreenSpotHFRepo.V2.value,
            system_prompt=system_prompt,
        )

    @classmethod
    def pro(cls, system_prompt: LocalizationPrompt) -> "ScreenSpotConfig":
        """Create a config for ScreenSpot pro test dataset"""
        return cls(
            dataset_version=ScreenSpotDatasetVersion.PRO.value,
            hf_repo=ScreenSpotHFRepo.PRO.value,
            system_prompt=system_prompt,
        )
