guidellm.utils.arg_string

class ArgStringParseError(Exception):
    """Exception raised for errors during argument string parsing."""

class ArgStringParser:
    """
    Decoder for converting argument strings to Python dictionaries.

    Parses comma-separated key=value strings with support for nested dictionaries
    (dot notation), list indexing (bracket notation), and automatic type conversion
    using YAML safe_load. Similar to json.JSONDecoder.

    Example:
    ::
        parser = ArgStringParser()
        result = parser.decode("prompt_tokens=10,output_tokens=30")
        # {'prompt_tokens': 10, 'output_tokens': 30}

        # Custom fill value for sparse lists
        parser = ArgStringParser(fill_value=lambda: 0)
        result = parser.decode("items[5]=value")
        # {'items': [0, 0, 0, 0, 0, 'value']}
    """

    def __init__(
        self,
        skip_invalid: bool = False,
        fill_value: Callable[[], Any] | None = None,
        allow_overwrite: bool = False,
        key_delimiter: str = ".",
        split_delimiter: str = ",",
    ):
        """
        Initialize the argument string parser.

        :param skip_invalid: If True, skips invalid key=value pairs rather than raising
        :param fill_value: Callable that returns fill value for sparse lists.
            Defaults to lambda: None.
        :param allow_overwrite: If True, allows overwriting existing values.
            Defaults to False, which raises an error on overwrites.
        """
        self.skip_invalid = skip_invalid
        self.fill_value = fill_value() if fill_value is not None else None
        self.allow_overwrite = allow_overwrite
        self.key_delimiter = key_delimiter
        self.split_delimiter = split_delimiter

    def decode(self, s: str) -> dict[str, Any]:
        """
        Decode an arg string into a Python dictionary with automatic type conversion.

        Converts comma-separated key=value pairs into nested dictionary structures.
        Supports dot notation for nested dictionaries (parent.child=value), bracket
        notation for lists (items[0]=value), and combined syntax.
        Values are parsed using YAML safe_load for automatic type conversion.

        :param s: Comma-separated list of key=value pairs supporting nested
            structures via dot and bracket notation
        :return: Parsed dictionary with appropriate types and nesting
        :raises ArgStringParseError: If the string format is invalid

        Example:
        ::
            parser = ArgStringParser()

            parser.decode("prompt_tokens=10,output_tokens=30")
            # {'prompt_tokens': 10, 'output_tokens': 30}

            parser.decode("prompt.tokens=15,output.tokens=45")
            # {'prompt': {'tokens': 15}, 'output': {'tokens': 45}}

            parser.decode("fruit[0]=apple,fruit[1]=orange")
            # {'fruit': ['apple', 'orange']}

            parser.decode("fruit[0]=apple,fruit[3]=orange")
            # {'fruit': ['apple', None, None, 'orange']}

            parser.decode("items[0].weight=10,items[0].count=5")
            # {'items': [{'weight': 10, 'count': 5}]}
        """
        result: dict[str, Any] = {}
        assigned_list_values: _ListValuePaths = set()

        if not s or not s.strip():
            return result

        # Split by commas
        pairs = s.split(self.split_delimiter)

        for idx, pair_raw in enumerate(pairs):
            pair = pair_raw.strip()
            # Handle invalid pairs
            if not pair or "=" not in pair:
                if self.skip_invalid:
                    continue
                raise ArgStringParseError(
                    f"Invalid key=value expression at {idx}: '{s}'"
                )

            key_raw, value_raw = pair.split("=", 1)
            key = key_raw.strip()
            value = value_raw.strip()

            # Handle empty keys
            if not key:
                if self.skip_invalid:
                    continue
                raise ArgStringParseError(
                    f"Empty key in key=value expression at {idx}: '{s}'"
                )

            # Parse the key into segments
            segments = self._parse_key(key)

            # Convert value to appropriate type
            parsed_value = yaml.safe_load(value)

            # Set the value in the nested structure
            self._value_set_nested(
                result,
                segments,
                parsed_value,
                assigned_list_values=assigned_list_values,
            )

        return result

    def set(self, obj: dict, path: str, value: Any) -> None:
        """
        Update a nested dictionary with a value at the specified path.

        This method allows updating a nested dictionary structure using a path
        string that supports dot notation for nested dictionaries and bracket
        notation for list indices. It creates intermediate dictionaries and lists
        as needed, filling sparse list gaps with the specified fill value.

        :param obj: The dictionary to update in place
        :param path: The path string indicating where to set the value, e.g.,
            "parent.child[0].name"
        :param value: The value to set at the specified path
        :raises ArgStringParseError: If the path format is invalid or if overwriting
            existing values is not allowed based on the parser's configuration
        """
        segments = self._parse_key(path)
        self._value_set_nested(obj, segments, value)

    def get(self, obj: dict, path: str) -> Any:
        """
        Retrieve a value from a nested dictionary using a path string.

        This method allows retrieving values from a nested dictionary structure
        using a path string that supports dot notation for nested dictionaries and
        bracket notation for list indices.

        :param obj: The dictionary to retrieve the value from
        :param path: The path string indicating where to get the value, e.g.,
            "parent.child[0].name"
        :return: The value at the specified path, or None if the path does not exist
        :raises ArgStringParseError: If the path format is invalid
        :raises KeyError: If a dictionary key in the path does not exist
        :raises IndexError: If a list index in the path is out of range
        """
        path_list = path.split(self.key_delimiter)
        segments = self._parse_key(path)
        current: Any = obj

        for i, (name, index) in enumerate(segments, 1):
            if not isinstance(current, dict) or name not in current:
                path_so_far = self.key_delimiter.join(path_list[:i])
                raise KeyError(f"Key '{name}' not found in '{path_so_far}'")
            current = current[name]

            if index is not None:
                if not isinstance(current, list) or index >= len(current):
                    path_so_far = self.key_delimiter.join(path_list[:i])
                    raise IndexError(f"Index '{path_so_far}' out of range")
                current = current[index]

        return current

    def _parse_key(self, key: str) -> list[tuple[str, int | None]]:
        """
        Parse a key string into segments with optional array indices.

        Splits a key by dots and identifies array access patterns using
        bracket notation. Returns a list of (name, index) tuples where index
        is None for dictionary keys and an integer for list indices.

        :param key: Key string to parse, may contain dots and brackets
        :return: List of (name, index) tuples representing the parsed key segments
        """
        segments: list[tuple[str, int | None]] = []

        # Split by dots first
        parts = key.split(self.key_delimiter)

        for part_raw in parts:
            # Check if part has array notation: name[index]
            part = part_raw.strip()
            match = re.match(r"^([^\[]+)\[(\d+)\]$", part)
            if match:
                name = match.group(1)
                index = int(match.group(2))
                segments.append((name, index))
            else:
                segments.append((part, None))

        return segments

    def _list_ensure_size(self, lst: list, index: int) -> None:
        """
        Ensure a list is large enough to access the given index.

        Extends the list with fill values until it can safely accommodate
        the specified index position.

        :param lst: The list to resize if necessary
        :param index: The minimum index that must be accessible
        """
        while len(lst) <= index:
            lst.append(self.fill_value)

    def _list_set_value(
        self,
        current: dict[str, Any],
        name: str,
        index: int,
        value: Any,
        *,
        assigned: bool = False,
    ) -> None:
        """
        Set a value in a list at the given index within a dictionary.

        Creates the list if it doesn't exist and ensures it's large enough
        to accommodate the specified index before setting the value.

        :param current: The current dictionary context containing the list
        :param name: The key name for the list in the dictionary
        :param index: The list index where the value should be set
        :param value: The value to set at the specified index
        :param assigned: Whether the list position was explicitly assigned
        :raises ArgStringParseError: If overwriting and allow_overwrite is False
        """
        if name not in current:
            current[name] = []
        elif not isinstance(current[name], list):
            # Trying to replace non-list with list
            if not self.allow_overwrite:
                raise ArgStringParseError(
                    f"Cannot overwrite non-list value at '{name}' with list"
                )
            current[name] = []

        self._list_ensure_size(current[name], index)

        # Check for overwrite of non-fill values
        existing_value = current[name][index]
        if not self.allow_overwrite and (assigned or existing_value != self.fill_value):
            raise ArgStringParseError(
                f"Cannot overwrite existing value at '{name}[{index}]'"
            )

        current[name][index] = value

    def _list_navigate_to_element(
        self,
        current: dict[str, Any],
        name: str,
        index: int,
        *,
        assigned: bool = False,
    ) -> Any:
        """
        Navigate to or create a list element for intermediate path segments.

        Creates the list if it doesn't exist, ensures it's large enough for the index,
        and initializes the element as an empty dictionary if needed. Returns the
        element for further navigation.

        :param current: The current dictionary context containing the list
        :param name: The key name for the list in the dictionary
        :param index: The list index to navigate to
        :param assigned: Whether the list position was explicitly assigned
        :return: The element at the given index (creating empty dict if needed)
        :raises ArgStringParseError: If overwriting and allow_overwrite is False
        """
        if name not in current:
            current[name] = []
        elif not isinstance(current[name], list):
            # Trying to replace non-list with list
            if not self.allow_overwrite:
                raise ArgStringParseError(
                    f"Cannot overwrite non-list value at '{name}' with list"
                )
            current[name] = []

        self._list_ensure_size(current[name], index)

        # Create nested structure if needed (only for fill values)
        existing_value = current[name][index]
        if existing_value == self.fill_value and not assigned:
            current[name][index] = {}
        elif not isinstance(existing_value, dict):
            # Trying to navigate into non-dict value
            if not self.allow_overwrite:
                raise ArgStringParseError(
                    f"Cannot overwrite non-dict value at '{name}[{index}]' with dict"
                )
            current[name][index] = {}

        return current[name][index]

    def _dict_navigate_to_key(
        self,
        current: dict[str, Any],
        name: str,
    ) -> dict[str, Any]:
        """
        Navigate to or create a dictionary key for intermediate path segments.

        Creates the dictionary key if it doesn't exist and returns the nested
        dictionary for further navigation.

        :param current: The current dictionary context
        :param name: The key name to navigate to
        :return: The nested dictionary at the given key
        :raises ArgStringParseError: If overwriting and allow_overwrite is False
        """
        if name not in current:
            current[name] = {}
        elif not isinstance(current[name], dict):
            # Trying to navigate into non-dict value
            if not self.allow_overwrite:
                raise ArgStringParseError(
                    f"Cannot overwrite non-dict value at '{name}' with dict"
                )
            current[name] = {}

        return current[name]

    def _value_set_nested(
        self,
        result: dict[str, Any],
        segments: list[tuple[str, int | None]],
        value: Any,
        *,
        assigned_list_values: _ListValuePaths | None = None,
    ) -> None:
        """
        Set a value in a nested dictionary and list structure.

        Navigates through the segment path, creating intermediate dictionaries and lists
        as needed. For sparse lists, fills gaps with fill values. Supports mixed nesting
        of dictionaries and lists.

        :param result: The root dictionary to modify in place
        :param segments: List of (name, index) tuples representing the navigation path
        :param value: The value to set at the final destination
        :param assigned_list_values: List positions explicitly assigned while decoding
        :raises ArgStringParseError: If overwriting and allow_overwrite is False
        """
        current: Any = result

        for i, (name, index) in enumerate(segments):
            is_last = i == len(segments) - 1
            path = tuple(segments[: i + 1])
            assigned = assigned_list_values is not None and path in assigned_list_values

            if is_last:
                # If we are at the last segment, set the value
                if index is not None:
                    self._list_set_value(
                        current,
                        name,
                        index,
                        value,
                        assigned=assigned,
                    )
                    if assigned_list_values is not None:
                        assigned_list_values.add(path)
                else:
                    if name in current and not self.allow_overwrite:
                        raise ArgStringParseError(
                            f"Cannot overwrite existing value at key '{name}'"
                        )
                    current[name] = value
            # If not last, navigate/create intermediate structures
            elif index is not None:
                # Intermediate segment with array access
                current = self._list_navigate_to_element(
                    current,
                    name,
                    index,
                    assigned=assigned,
                )
            else:
                # Intermediate segment with dict access
                current = self._dict_navigate_to_key(current, name)