guidellm.data.deserializers

class DataNotSupportedError(Exception):
    """
    Exception raised when the data format is not supported by deserializer or config.
    """

@DataArgs.register(["huggingface", "hf"])
class HuggingFaceDataArgs(DataArgs):
    """Model for Hugging Face dataset deserializer arguments."""

    kind: Literal["huggingface", "hf"] = Field(
        default="huggingface",
        description="Type identifier for the Hugging Face dataset deserializer.",
    )
    source: str | Dataset | IterableDataset | DatasetDict | IterableDatasetDict = Field(
        validation_alias=AliasChoices("source", "src", "from", "path", "name"),
        description=(
            "Data input for the Hugging Face dataset deserializer. This can be a "
            "Dataset, IterableDataset, DatasetDict, IterableDatasetDict, a string or "
            "Path to a local dataset directory or a local .py dataset script, or a "
            "dataset identifier from the Hugging Face Hub."
        ),
        examples=["wikimedia/structured-wikipedia", "./dataset.json"],
    )

@DatasetDeserializerFactory.register(["huggingface", "hf"])
class HuggingFaceDatasetDeserializer(DatasetDeserializer):
    """Deserializer for Hugging Face datasets."""

    def __call__(
        self,
        config: HuggingFaceDataArgs,
        processor_factory: Callable[[], PreTrainedTokenizerBase],
        random_seed: int,
    ) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
        _ = (processor_factory, random_seed)
        data = config.source

        if isinstance(
            data, Dataset | IterableDataset | DatasetDict | IterableDatasetDict
        ):
            return data

        load_error = None

        if (
            isinstance(data, str | Path)
            and (path := Path(data)).exists()
            and ((path.is_file() and path.suffix == ".py") or path.is_dir())
        ):
            # Handle python script or nested python script in a directory
            try:
                return load_dataset(str(data), **config.load_kwargs)
            except (
                FileNotFoundDatasetsError,
                DatasetNotFoundError,
                DataFilesNotFoundError,
            ) as err:
                load_error = err
            except Exception:  # noqa: BLE001
                # Try loading as a local dataset directory next
                try:
                    return load_from_disk(str(data), **config.load_kwargs)
                except (
                    FileNotFoundDatasetsError,
                    DatasetNotFoundError,
                    DataFilesNotFoundError,
                ) as err2:
                    load_error = err2

        try:
            # Handle dataset identifier from the Hugging Face Hub
            return load_dataset(str(data), **config.load_kwargs)
        except (
            FileNotFoundDatasetsError,
            DatasetNotFoundError,
            DataFilesNotFoundError,
        ) as err:
            load_error = err

        not_supported = DataNotSupportedError(
            "Unsupported data for HuggingFaceDatasetDeserializer, "
            "expected Dataset, IterableDataset, DatasetDict, IterableDatasetDict, "
            "str or Path to a local dataset directory or a local .py dataset script, "
            f"got {data} and HF load error: {load_error}"
        )

        if load_error is not None:
            raise not_supported from load_error
        else:
            raise not_supported

@DataArgs.register("in_memory_dict")
class InMemoryDictDataArgs(DataArgs):
    """Model for in-memory data deserializer arguments."""

    kind: Literal["in_memory_dict"] = Field(  # type: ignore[assignment]
        default="in_memory_dict",
        description="Type identifier for the in-memory data deserializer.",
    )
    data: dict[str, list] = Field(
        description="In-memory data input for the dataset deserializer.",
        examples=[{"column1": [1, 2, 3], "column2": [4, 5, 6]}],
    )

@DataArgs.register(_DESERIALIZER_TYPE)
class SyntheticImageDataArgs(SyntheticVisionDataArgs):
    """Model for synthetic image dataset deserializer arguments."""

    kind: Literal["synthetic_image"] = Field(  # type: ignore[assignment]
        default="synthetic_image",
        description="Type identifier for the synthetic image dataset configuration.",
    )
    width: int | None = Field(
        description="Image width in pixels.",
        gt=0,
        default=None,
    )
    height: int | None = Field(
        description="Image height in pixels.",
        gt=0,
        default=None,
    )
    resolution: str | None = Field(
        description="Resolution shortcut such as '720p' or '1080p'.",
        default=None,
    )
    aspect_ratio: str | None = Field(
        description="Aspect ratio override, e.g. '16:9' or '4:3'.",
        default=None,
    )
    format: Literal["jpeg", "png"] = Field(
        description="Encoded image format.",
        default="jpeg",
    )
    jpeg_quality: int = Field(
        description="JPEG quality 1..100. Ignored when format='png'.",
        ge=1,
        le=100,
        default=85,
    )
    content: Literal["gradient", "noise", "solid", "checkerboard"] = Field(
        description="Pixel content to synthesize.",
        default="gradient",
    )
    images_per_request: int = Field(
        description="Number of images per emitted row.",
        ge=1,
        default=1,
    )

    @model_validator(mode="after")
    def _resolve_dimensions(self) -> SyntheticImageDataArgs:
        w = self.width
        h = self.height
        if self.resolution is not None:
            preset = RESOLUTION_PRESETS.get(self.resolution.lower())
            if preset is None:
                raise ValueError(
                    f"Unknown resolution '{self.resolution}'. Known: "
                    f"{sorted(RESOLUTION_PRESETS)}"
                )
            preset_w, preset_h = preset
            if h is None:
                h = preset_h
            if w is None:
                w = (
                    int(round(h * parse_aspect_ratio(self.aspect_ratio)))
                    if self.aspect_ratio is not None
                    else preset_w
                )
        elif self.aspect_ratio is not None:
            if h is not None and w is None:
                w = int(round(h * parse_aspect_ratio(self.aspect_ratio)))
            elif w is not None and h is None:
                h = int(round(w / parse_aspect_ratio(self.aspect_ratio)))

        if w is None or h is None:
            raise ValueError(
                "synthetic_image config requires width and height, either "
                "explicitly or via resolution/aspect_ratio."
            )
        self.width = int(w) - (int(w) % 2)
        self.height = int(h) - (int(h) % 2)
        if self.width <= 0 or self.height <= 0:
            raise ValueError(
                f"Resolved image dims must be positive, got {self.width}x{self.height}"
            )
        return self

@DataArgs.register("synthetic_text")
class SyntheticTextDataArgs(DataArgs):
    """Model for synthetic text dataset deserializer arguments."""

    kind: Literal["synthetic_text"] = Field(  # type: ignore[assignment]
        default="synthetic_text",
        description="Type identifier for the synthetic text dataset configuration.",
    )
    prompt_tokens: int = Field(
        description="The average number of text tokens generated for prompts.",
        gt=0,
        examples=[30],
    )
    prompt_tokens_stdev: int | None = Field(
        description="The standard deviation of the tokens generated for prompts.",
        gt=0,
        default=None,
        examples=[3],
    )
    prompt_tokens_min: int | None = Field(
        description="The minimum number of text tokens generated for prompts.",
        gt=0,
        default=None,
        examples=[10],
    )
    prompt_tokens_max: int | None = Field(
        description="The maximum number of text tokens generated for prompts.",
        gt=0,
        default=None,
        examples=[30],
    )
    output_tokens: int | None = Field(
        description=(
            "The average number of text tokens generated for outputs. "
            "When omitted, output tokens are not sampled and ``max_tokens`` is left "
            "to the backend default. Useful for endpoints that do not produce "
            "output tokens (e.g. embeddings)."
        ),
        gt=0,
        default=None,
        examples=[10],
    )
    output_tokens_stdev: int | None = Field(
        description="The standard deviation of the tokens generated for outputs.",
        gt=0,
        default=None,
        examples=[3],
    )
    output_tokens_min: int | None = Field(
        description="The minimum number of text tokens generated for outputs.",
        gt=0,
        default=None,
        examples=[10],
    )
    output_tokens_max: int | None = Field(
        description="The maximum number of text tokens generated for outputs.",
        gt=0,
        default=None,
        examples=[30],
    )
    delay: float | None = Field(
        description='The average requeue delay, or "think time" for prompts.',
        gt=0,
        default=None,
        examples=[10.0],
    )
    delay_stdev: float | None = Field(
        description=(
            'The standard deviation of requeue delays, or "think time" for prompts.'
        ),
        gt=0,
        default=None,
        examples=[1.0],
    )
    delay_min: float | None = Field(
        description='The minimum requeue delay, or "think time" for prompts.',
        ge=0,
        default=None,
        examples=[0.5],
    )
    delay_max: float | None = Field(
        description='The maximum requeue delay, or "think time" for prompts.',
        gt=0,
        default=None,
        examples=[5.0],
    )
    turns: int = Field(
        description=(
            "The number of user turns in the conversation. "
            "Each tool-calling user turn automatically generates an additional "
            "tool_response_injection request, so the total request count per "
            "conversation is turns + len(tool_call_turns)."
        ),
        gt=0,
        default=1,
    )
    tool_call_turns: list[int] = Field(
        description=(
            "Which user turns should include tool definitions and expect "
            "tool-call responses. Indices are 0-based into the user turns "
            "(not the expanded request list). An int N means 'the first "
            "N user turns'; a list of ints specifies explicit indices "
            "(e.g. [0, 2]); -1 means all turns. Normalized to a sorted "
            "list after validation. "
            "When 0 or [] (default), no tool calling is configured."
        ),
        default_factory=list,
        examples=[1, [0, 1]],
    )
    tools: list[dict[str, Any]] | None = Field(
        description=(
            "Tool definitions in OpenAI format. When tool_call_turns is non-empty "
            "and this is None, a static placeholder tool definition is used."
        ),
        default=None,
        examples=[
            {
                "type": "function",
                "function": {
                    "name": "get_data",
                    "description": "Retrieve data from the system",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {"type": "string", "description": "The query"}
                        },
                        "required": ["query"],
                    },
                },
            }
        ],
    )
    tool_response_tokens: int | None = Field(
        description=(
            "Average number of tokens for synthetic tool call responses. "
            "When None (default), a short placeholder response is used."
        ),
        gt=0,
        default=None,
        examples=[10],
    )
    tool_response_tokens_stdev: int | None = Field(
        description="Standard deviation for tool response token count.",
        gt=0,
        default=None,
        examples=[1],
    )
    tool_response_tokens_min: int | None = Field(
        description="Minimum number of tokens for tool response.",
        gt=0,
        default=None,
        examples=[5],
    )
    tool_response_tokens_max: int | None = Field(
        description="Maximum number of tokens for tool response.",
        gt=0,
        default=None,
        examples=[20],
    )
    server_tool_call_turns: list[int] = Field(
        description=(
            "Which user turns use server-side tool calling. "
            "These turns are marked as server_tool_call so tool_choice='none' "
            "is not applied, letting the server use its configured tools. "
            "No injection turn is created. Must not overlap with "
            "tool_call_turns. Indices are 0-based into user turns. "
            "An int N means 'the first N user turns'; a list of ints "
            "specifies explicit indices (e.g. [0, 2]); -1 means all turns."
        ),
        default_factory=list,
    )

    prefix_buckets: list[SyntheticTextPrefixBucketConfig] | None = Field(
        description="Buckets for the prefix tokens distribution.",
        default=None,
        examples=[
            {"bucket_weight": 100, "prefix_count": 1, "prefix_tokens": 0},
        ],
    )

    @model_validator(mode="after")
    def check_prefix_options(self) -> SyntheticTextDataArgs:
        if self.__pydantic_extra__ is not None:
            prefix_count = self.__pydantic_extra__.get("prefix_count", None)  # type: ignore[attr-defined]
            prefix_tokens = self.__pydantic_extra__.get("prefix_tokens", None)  # type: ignore[attr-defined]

            if prefix_count is not None or prefix_tokens is not None:
                if self.prefix_buckets:
                    raise ValueError(
                        "prefix_buckets is mutually exclusive"
                        " with prefix_count and prefix_tokens"
                    )

                self.prefix_buckets = [
                    SyntheticTextPrefixBucketConfig(
                        prefix_count=prefix_count or 1,
                        prefix_tokens=prefix_tokens or 0,
                    )
                ]

        return self

    @field_validator("tool_call_turns", "server_tool_call_turns", mode="before")
    @classmethod
    def _coerce_tool_call_turns(
        cls, v: int | str | list[int], info: ValidationInfo
    ) -> list[int]:
        """Convert an int N to [0, ..., N-1]; pass lists through sorted.

        Strings are parsed as JSON to support CLI/env-var coercion.
        The value ``-1`` is converted to the sentinel ``[-1]`` which is
        expanded to all turn indices by :meth:`_validate_tool_call_turn_indices`
        once ``self.turns`` is available.
        """
        field = info.field_name
        if isinstance(v, str):
            try:
                v = json.loads(v)
            except (json.JSONDecodeError, ValueError) as err:
                raise ValueError(
                    f"{field} string must be a JSON int or list of ints, got {v!r}"
                ) from err
        if isinstance(v, int):
            if v == -1:
                return [-1]
            if v < 0:
                raise ValueError(f"{field} int must be >= 0 or -1 for all")
            return list(range(v))
        if not isinstance(v, list):
            raise ValueError(
                f"{field} must be int, list[int], or a JSON representation"
                f" of either, got {type(v)}"
            )
        if len(v) != len(set(v)):
            raise ValueError(f"{field} list must not contain duplicates")
        return sorted(v)

    @model_validator(mode="after")
    def _validate_tool_call_turn_indices(self) -> SyntheticTextDataArgs:
        """Ensure all tool call turn indices are within [0, turns) and don't overlap.

        The sentinel ``[-1]`` is expanded to ``list(range(self.turns))``
        before validation.
        """
        # Expand -1 sentinel ("all turns") for both fields
        if self.tool_call_turns == [-1]:
            self.tool_call_turns = list(range(self.turns))
        if self.server_tool_call_turns == [-1]:
            self.server_tool_call_turns = list(range(self.turns))

        for idx in self.tool_call_turns:
            if idx < 0 or idx >= self.turns:
                raise ValueError(
                    f"tool_call_turns index {idx} out of range [0, {self.turns})"
                )
        for idx in self.server_tool_call_turns:
            if idx < 0 or idx >= self.turns:
                raise ValueError(
                    f"server_tool_call_turns index {idx} out of range [0, {self.turns})"
                )
        overlap = set(self.tool_call_turns) & set(self.server_tool_call_turns)
        if overlap:
            raise ValueError(
                f"tool_call_turns and server_tool_call_turns must not overlap; "
                f"overlapping indices: {sorted(overlap)}"
            )
        return self

class SyntheticTextDataset(IterableDataset):
    def __init__(
        self,
        config: SyntheticTextDataArgs,
        processor: PreTrainedTokenizerBase,
        random_seed: int = 42,
    ):
        self.config = config
        self.processor = processor
        self.random_seed = random_seed

        # Create the examples iterable
        ex_iterable = _SyntheticTextExamplesIterable(
            config=config,
            processor=processor,
            random_seed=random_seed,
        )

        # Initialize parent with proper ex_iterable
        super().__init__(
            ex_iterable=ex_iterable,
            info=DatasetInfo(
                description="Synthetic text dataset generator",
                features=ex_iterable.features,
            ),
        )

    def set_epoch(self, epoch: int):
        """Set the epoch for the dataset iteration."""
        if isinstance(self._ex_iterable, _SyntheticTextExamplesIterable):
            self._ex_iterable.iteration_count = epoch

@DataArgs.register(_DESERIALIZER_TYPE)
class SyntheticVideoDataArgs(SyntheticVisionDataArgs):
    """Model for synthetic video dataset deserializer arguments."""

    kind: Literal["synthetic_video"] = Field(  # type: ignore[assignment]
        default="synthetic_video",
        description="Type identifier for the synthetic video dataset configuration.",
    )
    width: int | None = Field(
        description="Frame width in pixels.",
        gt=0,
        default=None,
    )
    height: int | None = Field(
        description="Frame height in pixels.",
        gt=0,
        default=None,
    )
    resolution: str | None = Field(
        description="Resolution shortcut such as '720p' or '1080p'.",
        default=None,
    )
    aspect_ratio: str | None = Field(
        description="Aspect ratio override, e.g. '16:9' or '4:3'.",
        default=None,
    )
    frames: int = Field(
        description="Number of frames in the clip.",
        ge=1,
    )
    fps: float = Field(
        description="Frames per second.",
        gt=0.0,
        default=1.0,
    )
    format: Literal["mp4"] = Field(
        description="Container / codec. Only mp4 (h264, yuv420p) in v1.",
        default="mp4",
    )
    video_bitrate: str | None = Field(
        description="Optional libx264 bitrate string, e.g. '500k'.",
        default=None,
    )
    content: Literal["gradient", "noise"] = Field(
        description="Frame content to synthesize.",
        default="gradient",
    )

    @model_validator(mode="after")
    def _resolve_dimensions(self) -> SyntheticVideoDataArgs:
        w = self.width
        h = self.height
        if self.resolution is not None:
            preset = RESOLUTION_PRESETS.get(self.resolution.lower())
            if preset is None:
                raise ValueError(
                    f"Unknown resolution '{self.resolution}'. Known: "
                    f"{sorted(RESOLUTION_PRESETS)}"
                )
            preset_w, preset_h = preset
            if h is None:
                h = preset_h
            if w is None:
                w = (
                    int(round(h * parse_aspect_ratio(self.aspect_ratio)))
                    if self.aspect_ratio is not None
                    else preset_w
                )
        elif self.aspect_ratio is not None:
            if h is not None and w is None:
                w = int(round(h * parse_aspect_ratio(self.aspect_ratio)))
            elif w is not None and h is None:
                h = int(round(w / parse_aspect_ratio(self.aspect_ratio)))

        if w is None or h is None:
            raise ValueError(
                "synthetic_video config requires width and height, either "
                "explicitly or via resolution/aspect_ratio."
            )
        self.width = int(w) - (int(w) % 2)
        self.height = int(h) - (int(h) % 2)
        if self.width <= 0 or self.height <= 0:
            raise ValueError(
                f"Resolved video dims must be positive, got {self.width}x{self.height}"
            )
        return self

class TraceDataArgs(DataArgs):
    """Abstract class meant to be inherited by a trace format.
    For testing, use `trace_minimal.MinimalTraceFormatArgs` instead."""

    kind: str = Field(
        description="Type identifier for the trace dataset deserializer.",
    )
    path: Path = Field(description="Path to the trace file.")
    timestamp_column: str = Field(
        default="timestamp",
        description="Column name for timestamps in the trace file.",
    )
    prompt_tokens_column: str = Field(
        default="input_length",
        description="Column name for prompt token counts in the trace file.",
    )
    output_tokens_column: str = Field(
        default="output_length",
        description="Column name for output token counts in the trace file.",
    )

@DatasetDeserializerFactory.register(["trace_synthetic"])
class TraceDatasetDeserializer(DatasetDeserializer):
    """Dataset deserializer for all trace formats."""

    def __call__(
        self,
        config: TraceDataArgs,
        processor_factory: Callable[[], PreTrainedTokenizerBase],
        random_seed: int = 42,
    ) -> IterableDataset:
        if not config.path.exists():
            raise DataNotSupportedError(f"Trace file not found: {config.path}")
        if not config.path.is_file():
            raise DataNotSupportedError(f"Trace path is not a file: {config.path}")
        return TraceDataset(config, processor_factory(), random_seed)

class TraceFormatBase(Protocol):
    def __init__(self) -> None: ...

    def required_columns(self, config) -> Features: ...

    def validate_row(self, config, row: dict) -> None:
        """Called within `trace_common.TraceExamplesIterable` on initialization,
        immediately after doing its own checks on the row."""

    def create_prompt(
        self,
        config,
        row: dict,
        processor: PreTrainedTokenizerBase,
        faker: Faker,
    ) -> str:
        """Called within `trace_common.TraceExamplesIterable` on each iteration.
        Returns a generated synthetic prompt."""