"""
Utilities for working with XLSX files.
"""

import re
from io import BytesIO
from typing import Iterable

import pandas as pd
from openpyxl import load_workbook
from openpyxl.cell.rich_text import CellRichText, TextBlock
from openpyxl.cell.text import InlineFont
from openpyxl.styles import PatternFill
from openpyxl.utils import get_column_letter

from ..utils.performance import no_garbage_collection


def format_cell_value(cell_value: str) -> CellRichText:
    """[REQUIRES OPENPYXL] This function formats an OpenPyxl cell value with HTML-like
    tags for bold, italic, and underline.

    Args:
        cell_value (str): The cell value to format.

    Returns:
        CellRichText: The formatted cell value, with rich text formatting. It
            can be used to set the value of an OpenPyxl cell.
    """
    pattern = re.compile(r"(</?[a-zA-Z]+>)([^<]*)")
    parts = pattern.findall(cell_value)

    # Initial state
    bold: bool = False
    italic: bool = False

    # Underline state: NOT BOOLEAN, but a string with the underline style. See
    # https://openpyxl.readthedocs.io/en/stable/api/openpyxl.styles.fonts.html#openpyxl.styles.fonts.Font.u
    underline: str | None = None

    text_blocks = []

    for tag, text in parts:
        tag = tag[1:-1]
        if tag == "b":
            bold = True
        elif tag == "/b":
            bold = False
        elif tag == "i":
            italic = True
        elif tag == "/i":
            italic = False
        elif tag == "u":
            underline = "single"
        elif tag == "/u":
            underline = None
        else:
            continue

        # Apply formatting only if there's text to apply it to
        if text:
            text_blocks.append(
                TextBlock(InlineFont(b=bold, i=italic, u=underline), text)
            )

    # Construct a CellRichText object with the formatted text blocks
    # Ensure there's at least one block, or just return the original cell value
    if text_blocks:
        return CellRichText(*text_blocks)
    return CellRichText(TextBlock(InlineFont(), cell_value))


def format_output_xlsx(xlsx: BytesIO) -> bytes:
    """Given an Excel file in BytesIO format, this function formats the file
    for presentation to end users.
    """

    wb = load_workbook(xlsx)

    orchard_orange_light = "ffba85"

    fill_color = PatternFill(
        start_color=orchard_orange_light,
        end_color=orchard_orange_light,
        fill_type="solid",
    )

    # Temporary disable garbage collection to speed up the process
    with no_garbage_collection():
        for ws in wb.worksheets:
            # Give the first row a background color
            for cell in ws[1]:
                cell.fill = fill_color

            for column in ws.columns:
                column_autosize(ws, column)

                format_cache = {}

                # Recognize and apply HTML formatting to cells
                for cell in column:
                    cell_value = str(cell.value) if cell.value is not None else ""
                    cached = format_cache.get(cell_value, None)
                    if cached is None:
                        formatted_value = format_cell_value(cell_value)
                        format_cache[cell_value] = formatted_value
                        cell.value = formatted_value
                    else:
                        cell.value = cached

    return workbook_to_xlsx_bytes(wb)


def column_autosize(ws, column):
    # Uses a cache to avoid recalculating the same cell length multiple times when
    # the same value is repeated in the same column.
    value_cache = {}

    max_length = 0

    for cell in column:
        cached = value_cache.get(cell.value)
        if cached is None:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(str(cell.value))
            except Exception:  # noqa E722
                pass
            value_cache[cell.value] = max_length
        else:
            max_length = cached
    new_width = max_length + 2
    ws.column_dimensions[get_column_letter(column[0].column)].width = new_width


def df_to_xlsx_bytes(sheets: Iterable[tuple[str, pd.DataFrame]]) -> BytesIO:
    """Given a list of tuples with sheet names and DataFrames, this function
    creates an Excel file in BytesIO format.

    The order of the sheets is preserved in the output file.
    """
    output = BytesIO()
    with pd.ExcelWriter(output, engine="openpyxl") as writer:
        for sheet_name, df in sheets:
            df.to_excel(writer, sheet_name=sheet_name, index=False)
    output.seek(0)
    return output


def workbook_to_xlsx_bytes(wb) -> bytes:
    output = BytesIO()
    wb.save(output)
    output.seek(0)
    return output.getvalue()