from typing import Literal

from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmarkConfig
from screensuite.benchmarks.perception.websrc.prompt import WebSrcPrompt


class WebSrcConfig(HubBaseBenchmarkConfig):
    hf_repo: Literal["rootsautomation/websrc"]
    """HF repo name"""

    revision: str = "main"

    split: Literal["train", "dev"]
    """HF split name"""

    system_prompt: WebSrcPrompt = WebSrcPrompt.SHORT_PROMPT
    """System prompt"""

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

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

    @classmethod
    def dev(cls) -> "WebSrcConfig":
        """Create a config for WebSrc dev dataset"""
        return cls(
            hf_repo="rootsautomation/websrc",
            system_prompt=WebSrcPrompt.SHORT_PROMPT,
            split="dev",
        )
