# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import hashlib
import logging
import os
import re
from typing import Optional

from oumi.core.tokenizers.base_tokenizer import BaseTokenizer



[docs]
def sanitize_run_name(run_name: Optional[str]) -> Optional[str]:
    """Computes a sanitized version of wandb run name.

    A valid run name may only contain alphanumeric characters, dashes, underscores,
    and dots, with length not exceeding max limit.

    Args:
        run_name: The original raw value of run name.
    """
    if not run_name:
        return run_name

    # Technically, the limit is 128 chars, but we limit to 100 characters
    # because the system may generate aux artifact names e.g., by prepending a prefix
    # (e.g., "model-") to our original run name, which are also subject
    # to max 128 chars limit.
    _MAX_RUN_NAME_LENGTH = 100

    # Replace all unsupported characters with '_'.
    result = re.sub("[^a-zA-Z0-9\\_\\-\\.]", "_", run_name)
    if len(result) > _MAX_RUN_NAME_LENGTH:
        suffix = "..." + hashlib.shake_128(run_name.encode("utf-8")).hexdigest(8)
        result = result[0 : (_MAX_RUN_NAME_LENGTH - len(suffix))] + suffix

    if result != run_name:
        logger = logging.getLogger("oumi")
        logger.warning(f"Run name '{run_name}' got sanitized to '{result}'")
    return result




[docs]
def try_str_to_bool(s: str) -> Optional[bool]:
    """Attempts to convert a string representation to a boolean value.

    This function interprets various string inputs as boolean values.
    It is case-insensitive and recognizes common boolean representations.

    Args:
        s: The string to convert to a boolean.

    Returns:
        bool: The boolean interpretation of the input string, or `None`
            for unrecognized string values.

    Examples:
        >>> str_to_bool("true") # doctest: +SKIP
        True
        >>> str_to_bool("FALSE") # doctest: +SKIP
        False
        >>> str_to_bool("1") # doctest: +SKIP
        True
        >>> str_to_bool("no") # doctest: +SKIP
        False
        >>> str_to_bool("peach") # doctest: +SKIP
        None
    """
    s = s.strip().lower()

    if s in ("true", "yes", "1", "on", "t", "y"):
        return True
    elif s in ("false", "no", "0", "off", "f", "n"):
        return False
    return None




[docs]
def str_to_bool(s: str) -> bool:
    """Convert a string representation to a boolean value.

    This function interprets various string inputs as boolean values.
    It is case-insensitive and recognizes common boolean representations.

    Args:
        s: The string to convert to a boolean.

    Returns:
        bool: The boolean interpretation of the input string.

    Raises:
        ValueError: If the input string cannot be interpreted as a boolean.

    Examples:
        >>> str_to_bool("true") # doctest: +SKIP
        True
        >>> str_to_bool("FALSE") # doctest: +SKIP
        False
        >>> str_to_bool("1") # doctest: +SKIP
        True
        >>> str_to_bool("no") # doctest: +SKIP
        False
    """
    result = try_str_to_bool(s)

    if result is None:
        raise ValueError(f"Cannot convert '{s}' to boolean.")
    return result




[docs]
def compute_utf8_len(s: str) -> int:
    """Computes string length in UTF-8 bytes."""
    # This is inefficient: allocates a temporary copy of string content.
    # FIXME Can we do better?
    return len(s.encode("utf-8"))




[docs]
def get_editable_install_override_env_var() -> bool:
    """Returns whether OUMI_FORCE_EDITABLE_INSTALL env var is set to a truthy value."""
    s = os.environ.get("OUMI_FORCE_EDITABLE_INSTALL", "")
    mode = s.lower().strip()
    bool_result = try_str_to_bool(mode)
    if bool_result is not None:
        return bool_result
    return False



# Experimental function, only for developer usage.

[docs]
def set_oumi_install_editable(setup: str) -> str:
    """Tries to replace oumi PyPi installs with editable installation from source.

    For example, the following line:
        `pip install uv && uv pip -q install oumi[gpu,dev] vllm`
    will be replaced with:
        `pip install uv && uv pip -q install -e '.[gpu,dev]' vllm`

    Args:
        setup (str): The bash setup script to modify. May be multi-line.

    Returns:
        The modified setup script.
    """
    setup_lines = setup.split("\n")
    for i, line in enumerate(setup_lines):
        # Skip comments.
        if line.strip().startswith("#"):
            continue

        # In summary, this regex looks for variants of `pip install oumi` and replaces
        # the oumi package with an editable install from the current directory.
        #
        # Tip: Use https://regexr.com/ or an LLM to help understand the regex.
        # It captures any misc. tokens like flags for the pip and
        # install commands, in addition to any optional dependencies oumi is installed
        # with.
        #
        # `((?:[-'\"\w]+ +)*)` matches whitespace-separated tokens potentially
        # containing quotes, such as flag names and values.
        # `((?:[-'\",\[\]\w]+ +)*)` does the same, with the addition of commas and
        # brackets, which may be present for packages with optional dependencies.
        # Since these don't include special characters like && and ;, it shouldn't span
        # across multiple pip install commands.
        # `(?<!-e )` means we don't match if the previous token is -e. This means an
        # editable install of a local dir called "oumi" is being done, so we skip it.
        # NOTE: We ideally should check for `--editable` as well, but Python re doesn't
        # support lookbehinds with variable length.
        # We additionally consume quotation marks around oumi if present.
        # Finally, `(\[[^\]]*\])?['\"]?` captures optional dependencies, if present.
        pattern = (
            r"pip3? +((?:[-'\"\w]+ +)*)install +((?:[-'\",\[\]\w]+ +)*)"
            r"(?<!-e )['\"]?oumi(\[[^\]]*\])?['\"]?"
        )
        # Compared to the pattern we captured, the changes are replacing `oumi` with
        # `.` and adding `-e` to make the install editable.
        replacement = r"pip \1install \2-e '.\3'"

        result = re.sub(pattern, replacement, line)
        if result == line:
            continue
        # Replace the line in the setup script.
        logger = logging.getLogger("oumi")
        logger.info(f"Detected the following oumi installation: `{line}`")
        logger.info(f"Replaced with: `{result}`")
        setup_lines[i] = result
    return "\n".join(setup_lines)




[docs]
def truncate_to_max_tokens_limit(
    text: str,
    tokenizer: BaseTokenizer,
    *,
    max_tokens: int,
    truncation_side: str = "right",
) -> tuple[str, int]:
    """Truncates text to `max_length` in tokens.

    Args:
        text: A text prompt.
        tokenizer: The tokenizer used for encoding the data.
        max_tokens: Maximum number of tokens to keep.
        truncation_side: The side to truncate the tokens ("right" or "left").

    Returns:
        A tuple containing truncated text prompt and the number of tokens.
    """
    if max_tokens <= 0:
        raise ValueError("`max_tokens` must be a positive integer")
    elif truncation_side not in ("left", "right"):
        raise ValueError(
            f"Invalid truncation_side: '{truncation_side}'. Expected 'left' or 'right'."
        )

    if not text:
        return ("", 0)

    left_side = truncation_side == "left"

    # The `truncation_side` parameter isn't universally supported by all tokenizers.
    # Let's do left-side truncation as post-processing.
    result = tokenizer(
        text,
        return_offsets_mapping=True,
        return_length=True,
        max_length=(None if left_side else max_tokens),
        truncation=(not left_side),  # Left-side truncation is done as post-processing.
    )
    if "offset_mapping" not in result:
        raise RuntimeError(
            f"Tokenizer must return offset mapping for truncation! Got: {result.keys()}"
        )
    token2char_offsets = result["offset_mapping"]
    if not isinstance(token2char_offsets, list):
        raise RuntimeError(
            "offset_mapping returned by tokenizer is not a list! "
            f"Got: {type(token2char_offsets)}"
        )

    truncated_text: str = ""
    num_truncated_tokens: int = 0
    if len(token2char_offsets) > 0:
        num_truncated_tokens = min(len(token2char_offsets), max_tokens)
        if left_side:
            lead_token_start = token2char_offsets[-num_truncated_tokens][0]
            if not (lead_token_start >= 0 and lead_token_start < len(text)):
                raise RuntimeError(
                    f"Truncation error: lead_token_start={lead_token_start} "
                    f"for text of length {len(text)}"
                )
            truncated_text = text[lead_token_start:]
        else:
            last_token_end = token2char_offsets[num_truncated_tokens - 1][1]
            if not (last_token_end >= 0 and last_token_end <= len(text)):
                raise RuntimeError(
                    f"Truncation error: last_token_end={last_token_end} "
                    f"for text of length {len(text)}"
                )
            truncated_text = text[:last_token_end]

    return (truncated_text, num_truncated_tokens)




[docs]
def truncate_text_pieces_to_max_tokens_limit(
    text_pieces: list[str],
    tokenizer: BaseTokenizer,
    *,
    max_tokens: int,
    truncation_side: str = "right",
) -> list[str]:
    """Truncates text pieces to total length not exceeding `max_length`.

    Args:
        text_pieces: A list of text prompts.
        tokenizer: The tokenizer used for encoding the data.
        max_tokens: Maximum number of tokens to keep in all text pieces combined.
        truncation_side: The side to truncate the tokens ("right" or "left").

    Returns:
        A list of truncated text prompts.
    """
    if max_tokens <= 0:
        raise ValueError("`max_tokens` must be a positive integer")
    remaining_tokens = max_tokens

    result = copy.deepcopy(text_pieces)
    if truncation_side == "left":
        result.reverse()

    for idx, text_piece in enumerate(result):
        if len(text_piece) == 0:
            continue
        elif remaining_tokens > 0:
            truncated_text_piece, num_tokens = truncate_to_max_tokens_limit(
                text_piece,
                tokenizer=tokenizer,
                max_tokens=remaining_tokens,
                truncation_side=truncation_side,
            )
            result[idx] = truncated_text_piece
            remaining_tokens -= num_tokens
        else:
            result[idx] = ""

    if truncation_side == "left":
        result.reverse()

    return result