import re

from .agents.types import TokenUsage


def format_diff_with_line_numbers(diff_text: str) -> str:
    """Parses a standard git diff output and adds old/new line numbers to each line, mimicking the output of scripts like git-diffn.

    Args:
        diff_text: A string containing the raw output from `git diff`.

    Returns:
        A string with the formatted diff, including line numbers.
    """
    output_lines = []

    # Regex to robustly parse the hunk header (e.g., "@@ -337,11 +337,13 @@")
    # It captures the starting line numbers for both the old and new files.
    hunk_header_pattern = re.compile(r"^@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@")

    old_line_num = 0
    new_line_num = 0

    # Ensure diff_text is a string before splitting
    if not isinstance(diff_text, str):
        return ""  # Or raise an error, depending on desired behavior

    for line in diff_text.splitlines():
        match = hunk_header_pattern.match(line)
        if match:
            # When a hunk header is found, reset the line counters
            # to the starting numbers specified in the header.
            old_line_num = int(match.group(1))
            new_line_num = int(match.group(2))
            output_lines.append(line)
            continue

        if not line:
            output_lines.append("")
            continue

        prefix = line[0]
        content = line[1:]

        # Based on the line prefix (+, -, or space), format and print the line.
        if prefix == "+":
            # This is an added line. It only has a line number in the new file.
            # Pad for alignment.
            formatted_line = f"      {new_line_num:<5}:+{content}"
            new_line_num += 1
        elif prefix == "-":
            # This is a deleted line. It only has a line number in the old file.
            # Pad for alignment.
            formatted_line = f"{old_line_num:<5}     :-{content}"
            old_line_num += 1
        elif prefix == " ":
            # This is a context line, existing in both files.
            formatted_line = f"{old_line_num:<5}, {new_line_num:<5}: {content}"
            old_line_num += 1
            new_line_num += 1
        else:
            # This line is not part of a hunk (e.g., "diff --git...", "index...").
            # Print it as-is.
            formatted_line = line

        output_lines.append(formatted_line)

    return "\n".join(output_lines)


def aggregate_token_usage(base: TokenUsage, additional: TokenUsage) -> TokenUsage:
    """Aggregate two TokenUsage objects by summing their values.

    This is a convenience wrapper around TokenUsage.__add__.
    You can also use: total = base + additional

    Args:
        base: Base token usage to add to
        additional: Additional token usage to add

    Returns:
        New TokenUsage with aggregated values
    """
    return base + additional
