"""Utility functions for working with trajectory data."""

import json
import re


def format_trajectory_json(data: dict) -> str:
    """Format trajectory JSON with compact numeric arrays on single lines.

    This formats the JSON with regular indentation but keeps large numeric
    arrays (like prompt_token_ids, completion_token_ids, logprobs) on a single line.

    Args:
        data: Dictionary representation of trajectory data

    Returns:
        Formatted JSON string with compact numeric arrays on single lines
    """
    # First, dump with standard formatting
    json_str = json.dumps(data, indent=2)

    # Compact arrays of numbers: put all elements on a single line
    def compact_numeric_array(match):
        full_match = match.group(0)
        numbers = re.findall(r'-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?', full_match)

        if not numbers:
            return full_match

        result = "[" + ", ".join(numbers) + "]"
        return result

    # Match arrays that span multiple lines with numbers (one per line)
    # Pattern: [ followed by whitespace/numbers/commas, ending with ]
    pattern = r'\[\s*\n\s*-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?(?:\s*,\s*\n\s*-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?)*\s*\n\s*\]'
    json_str = re.sub(pattern, compact_numeric_array, json_str, flags=re.MULTILINE)

    return json_str
