import xlsxwriter
import itertools
from typing import Iterable


def auto_fit_columns(
    worksheet: xlsxwriter.workbook.Worksheet,
    data: Iterable[list[str | int | float]],
) -> None:
    """Auto-fit column widths based on string lengths in a 2D list of data
    (including headers).

    Args:
        worksheet (xlsxwriter.workbook.Worksheet): The worksheet to set column
         widths for.
        data (Iterable[list[str | int | float]]): The 2D list of data to analyze. As
         xlsxwriter is write-only, we cannot read the data back from the worksheet
         and need to pass the data in. It can also be a generator.
    """
    data, data_copy = itertools.tee(data)
    try:
        first_row = next(data_copy)
    except StopIteration:
        return

    num_cols = len(first_row)
    col_widths = [0] * num_cols

    for row in data:
        for i in range(num_cols):
            if i < len(row):
                col_widths[i] = max(col_widths[i], len(str(row[i])))

    for i, width in enumerate(col_widths):
        worksheet.set_column(i, i, width + 2)