pdf/gateway/app/services/convert/formatters/xlsx_formatter.py

"""XLSX formatter from IDM."""

from __future__ import annotations

import io
import re
from datetime import date

from openpyxl import Workbook
from openpyxl.styles import Font
from openpyxl.utils import get_column_letter

from app.services.convert.formatters.xml_sanitize import sanitize_ooxml_text
from app.services.convert.idm.model import BlockType, Document
from app.services.convert.text.arabic_logical import to_logical_cell
from app.services.convert.text.numbers import (
    excel_number_format,
    parse_iso_date,
    parse_number,
)


def _block_visual(block) -> bool:
    """True when a block's text is still in visual order and needs repairing.

    Ordering is now resolved at extraction time — ``layout.glyphs`` and the OCR
    adapter sort each line by its own reading direction — so a block's source
    no longer implies visual order. Only an explicit span flag does, which is
    set when a producer genuinely could not determine direction.
    """
    return any(getattr(s, "visual_order", False) for s in (getattr(block, "spans", None) or []))


# Excel grid limits (xlsx). Exceeding either raises inside openpyxl.
MAX_XLSX_ROWS = 1_048_576
MAX_XLSX_COLS = 16_384


def _safe_sheet_name(name: str, used: set[str]) -> str:
    cleaned = re.sub(r"[\[\]\*\/\\\?\:]", "_", sanitize_ooxml_text(name))[:31] or "Sheet"
    base = cleaned
    n = 1
    while cleaned in used:
        suffix = f"_{n}"
        cleaned = base[: 31 - len(suffix)] + suffix
        n += 1
    used.add(cleaned)
    return cleaned


def coerce_cell(value: str, *, visual: bool = False) -> tuple[object, str | None]:
    """Return ``(python value, excel number format)`` for a cell string.

    Numbers are written as numbers so ``=SUM()`` works — storing them as text
    is the most common complaint about PDF-to-Excel output. The display format
    travels with the value so "4.50", "1,234" and "12%" still look the same in
    Excel while being numeric underneath.

    Labels that merely contain digits ("R1", "Q3", "00123") stay text: coercing
    them would rewrite row identifiers into values.
    """
    raw = sanitize_ooxml_text(to_logical_cell(value, visual=visual)).strip()
    if not raw:
        return "", None

    iso = parse_iso_date(raw)
    if iso:
        return date(*iso), "yyyy-mm-dd"

    parsed = parse_number(raw)
    if parsed is None:
        return raw, None

    fmt = excel_number_format(parsed)
    if parsed.is_integral and not parsed.percent and not parsed.currency:
        return int(parsed.value), fmt
    return parsed.value, fmt


def _coerce_cell(value: str, *, visual: bool = False):
    """Backwards-compatible single-value form used by existing callers/tests."""
    return coerce_cell(value, visual=visual)[0]


def _write_prose(ws, blocks, start_row: int = 1) -> int:
    r_idx = start_row
    for block in sorted(blocks, key=lambda b: b.reading_order):
        if block.type in (BlockType.header, BlockType.footer, BlockType.table, BlockType.figure):
            continue
        text = sanitize_ooxml_text(
            to_logical_cell(block.plain_text(), visual=_block_visual(block))
        ).strip()
        if not text:
            continue
        value, number_format = coerce_cell(text)
        cell = ws.cell(row=r_idx, column=1, value=value)
        if number_format:
            cell.number_format = number_format
        if block.type == BlockType.heading or any(s.bold for s in block.spans):
            cell.font = Font(bold=True)
        r_idx += 1
    return r_idx


def format_xlsx(doc: Document) -> tuple[bytes, list[str]]:
    warnings: list[str] = []
    wb = Workbook()
    default = wb.active
    any_table = False
    used_names: set[str] = set()
    first_sheet = True

    all_blocks = [b for page in doc.pages for b in page.blocks]
    high_conf_tables = [
        b
        for b in all_blocks
        if b.type == BlockType.table and b.cells and b.table_confidence >= 0.65
    ]

    # Always keep a Content sheet with non-table prose (invoice KV, etc.)
    content_needed = any(
        b.type not in (BlockType.table, BlockType.header, BlockType.footer, BlockType.figure)
        and b.plain_text().strip()
        for b in all_blocks
    )

    if len(high_conf_tables) > 1:
        for ti, table_block in enumerate(high_conf_tables):
            any_table = True
            name = _safe_sheet_name(f"Table{ti + 1}", used_names)
            if first_sheet:
                ws = default
                ws.title = name
                first_sheet = False
            else:
                ws = wb.create_sheet(name)
            _write_table(
                ws,
                table_block.cells,
                table_block.merges,
                freeze=table_block.table_confidence >= 0.75,
                warnings=warnings,
                visual=_block_visual(table_block),
            )
        if content_needed:
            name = _safe_sheet_name("Content", used_names)
            ws = wb.create_sheet(name)
            _write_prose(ws, all_blocks)
    else:
        for page in doc.pages:
            name = _safe_sheet_name(f"Page{page.index + 1}", used_names)
            if first_sheet:
                ws = default
                ws.title = name
                first_sheet = False
            else:
                ws = wb.create_sheet(name)

            tables = [b for b in page.blocks if b.type == BlockType.table and b.cells]
            prose_blocks = [b for b in page.blocks if b.type != BlockType.table]
            r_next = 1
            if prose_blocks:
                r_next = _write_prose(ws, prose_blocks, start_row=1)
                if tables:
                    r_next += 1
            if tables:
                any_table = True
                # Write first table starting at r_next
                _write_table_at(
                    ws,
                    tables[0].cells,
                    tables[0].merges,
                    start_row=r_next,
                    freeze=tables[0].table_confidence >= 0.75,
                    warnings=warnings,
                    visual=_block_visual(tables[0]),
                )
                col_offset = max(len(r) for r in tables[0].cells) + 2
                for extra in tables[1:]:
                    _write_table(
                        ws,
                        extra.cells,
                        extra.merges,
                        col_offset=col_offset,
                        warnings=warnings,
                        visual=_block_visual(extra),
                    )
                    col_offset += max(len(r) for r in extra.cells) + 2

    if not any_table:
        warnings.append("No table structure detected; exported paragraphs as a single column.")

    buf = io.BytesIO()
    wb.save(buf)
    return buf.getvalue(), warnings


def _write_table(
    ws,
    cells: list[list[str]],
    merges: list[tuple[int, int, int, int]] | None = None,
    *,
    col_offset: int = 0,
    freeze: bool = False,
    warnings: list[str] | None = None,
    visual: bool = False,
) -> None:
    _write_table_at(
        ws,
        cells,
        merges,
        start_row=1,
        col_offset=col_offset,
        freeze=freeze,
        warnings=warnings,
        visual=visual,
    )


def _write_table_at(
    ws,
    cells: list[list[str]],
    merges: list[tuple[int, int, int, int]] | None = None,
    *,
    start_row: int = 1,
    col_offset: int = 0,
    freeze: bool = False,
    warnings: list[str] | None = None,
    visual: bool = False,
) -> None:
    for r_idx, row in enumerate(cells):
        if start_row + r_idx > MAX_XLSX_ROWS:
            if warnings is not None:
                warnings.append(
                    f"Table truncated at {MAX_XLSX_ROWS} rows (Excel sheet limit)."
                )
            break
        for c_idx, cell in enumerate(row):
            if col_offset + c_idx + 1 > MAX_XLSX_COLS:
                if warnings is not None:
                    warnings.append(
                        f"Table truncated at {MAX_XLSX_COLS} columns (Excel sheet limit)."
                    )
                break
            number_format: str | None = None
            if isinstance(cell, str):
                val, number_format = coerce_cell(cell, visual=visual)
            elif cell is None:
                val = ""
            else:
                val = cell
            xl = ws.cell(row=start_row + r_idx, column=col_offset + c_idx + 1, value=val)
            if number_format:
                xl.number_format = number_format
            if r_idx == 0:
                xl.font = Font(bold=True)
            # Column width heuristic
            try:
                letter = get_column_letter(col_offset + c_idx + 1)
                cur = ws.column_dimensions[letter].width or 8
                ws.column_dimensions[letter].width = max(cur, min(28, len(str(val)) + 2))
            except Exception:
                pass
    for r0, c0, r1, c1 in merges or []:
        try:
            start = f"{get_column_letter(col_offset + c0 + 1)}{start_row + r0}"
            end = f"{get_column_letter(col_offset + c1 + 1)}{start_row + r1}"
            ws.merge_cells(f"{start}:{end}")
        except Exception as exc:
            if warnings is not None:
                warnings.append(f"XLSX merge failed ({r0},{c0})-({r1},{c1}): {exc}")
    if freeze and start_row >= 1:
        try:
            ws.freeze_panes = f"A{start_row + 1}"
        except Exception:
            pass