digest_builder.py

"""
digest_builder.py — Extracts a token-efficient workbook digest for AI analysis.

Does not call the AI. Returns a WorkbookDigest containing per-sheet summaries,
named ranges, and VBA metadata. The digest is consumed by ai_analyser.py.
"""

from __future__ import annotations

import logging
import re
from pathlib import Path

import openpyxl

import settings
from models import SheetSummary, WorkbookDigest

logger = logging.getLogger(__name__)

# Regex to find sheet references in formula strings: SheetName! or 'Sheet Name'!
_SHEET_REF_PATTERN = re.compile(r"'([^']+)'!|\b([A-Za-z0-9_]+)!")


def _formula_complexity(formula: str) -> int:
    """Compute a complexity score for a formula string."""
    return len(formula) + formula.count("(") * 5


def _extract_headers(ws) -> list[str]:
    """
    Collect string values from row 1 and column A, deduplicated and capped.
    Returns up to AI_MAX_HEADER_VALUES unique non-empty strings.
    """
    seen: set[str] = set()
    headers: list[str] = []

    # Row 1
    try:
        for cell in ws[1]:
            if cell.value and isinstance(cell.value, str):
                val = cell.value.strip()
                if val and val not in seen:
                    seen.add(val)
                    headers.append(val)
                    if len(headers) >= settings.AI_MAX_HEADER_VALUES:
                        return headers
    except Exception:
        pass

    # Column A (skip row 1 — already covered)
    try:
        for cell in ws["A"]:
            if cell.row == 1:
                continue
            if cell.value and isinstance(cell.value, str):
                val = cell.value.strip()
                if val and val not in seen:
                    seen.add(val)
                    headers.append(val)
                    if len(headers) >= settings.AI_MAX_HEADER_VALUES:
                        return headers
    except Exception:
        pass

    return headers


def _extract_sheet_references(formula: str, current_sheet: str) -> list[str]:
    """Parse a formula string for cross-sheet references (SheetName!)."""
    refs: set[str] = set()
    for m in _SHEET_REF_PATTERN.finditer(formula):
        name = m.group(1) or m.group(2)
        if name and name != current_sheet:
            refs.add(name)
    return list(refs)


def _extract_sheet_summary(
    ws,
    sensitive_sheets: set[str],
    sheet_named_ranges: dict[str, list[str]],
) -> SheetSummary:
    """Build a SheetSummary for a single worksheet."""
    name = ws.title
    is_sensitive = name in sensitive_sheets

    # Dimensions
    row_count = ws.max_row or 0
    col_count = ws.max_column or 0

    # Headers
    headers = _extract_headers(ws)

    # Top formulas by complexity
    formula_cells: list[tuple[str, str]] = []  # (cell_address, formula)
    referenced: set[str] = set()

    try:
        for row in ws.iter_rows():
            for cell in row:
                val = cell.value
                if isinstance(val, str) and val.startswith("="):
                    formula_cells.append((cell.coordinate, val))
                    for ref in _extract_sheet_references(val, name):
                        referenced.add(ref)
    except Exception:
        pass

    formula_cells.sort(key=lambda fc: _formula_complexity(fc[1]), reverse=True)
    top_formulas = [
        f"{addr}: {formula[:settings.AI_MAX_FORMULA_CHARS]}"
        + ("..." if len(formula) > settings.AI_MAX_FORMULA_CHARS else "")
        for addr, formula in formula_cells[: settings.AI_TOP_FORMULAS_PER_SHEET]
    ]

    named_ranges = sheet_named_ranges.get(name, [])

    return SheetSummary(
        name=name,
        row_count=row_count,
        col_count=col_count,
        is_sensitive=is_sensitive,
        headers=headers,
        top_formulas=top_formulas,
        named_ranges=named_ranges,
        references_sheets=sorted(referenced),
    )


def _extract_named_ranges(wb) -> tuple[dict[str, str], dict[str, list[str]]]:
    """
    Return (workbook_named_ranges, sheet_named_ranges).

    workbook_named_ranges: name -> formula/value string (workbook-scoped)
    sheet_named_ranges:    sheet_name -> [range_name, ...]
    """
    wb_ranges: dict[str, str] = {}
    sheet_ranges: dict[str, list[str]] = {}

    try:
        # openpyxl 3.1.x: wb.defined_names is a DefinedNameList (not a dict);
        # iterate directly to get DefinedName objects, each with a .name attribute.
        for defn in wb.defined_names:
            try:
                name = defn.name
                destinations = list(defn.destinations)
            except Exception:
                continue

            if destinations:
                # Sheet-scoped range
                sheet_name, coord = destinations[0]
                wb_ranges[name] = f"{sheet_name}!{coord}"
                sheet_ranges.setdefault(sheet_name, []).append(name)
            else:
                # Workbook-scoped or scalar
                try:
                    wb_ranges[name] = str(defn.attr_text or defn.value or "")
                except Exception:
                    wb_ranges[name] = ""
    except Exception:
        pass

    return wb_ranges, sheet_ranges


def _extract_vba_module_names(wb_path: Path, result) -> list[str]:
    """
    Return VBA module names by inspecting existing VBA findings if available.
    Never re-extracts VBA source.
    """
    if result is None:
        return []
    try:
        module_names: set[str] = set()
        for finding in result.findings:
            if finding.category == "VBA" and finding.check_id == 15:
                # The VBA presence finding description may list module names
                desc = finding.description
                # Extract quoted module names from description text
                for m in re.findall(r"'([^']+)'", desc):
                    module_names.add(m)
        return sorted(module_names)
    except Exception:
        return []


def build_digest(wb_path: Path, workbook_analyser_result=None) -> WorkbookDigest:
    """
    Build and return a WorkbookDigest from the workbook at wb_path.

    Args:
        wb_path:                  Path to the .xlsx or .xlsm file.
        workbook_analyser_result: Optional WorkbookAnalysisResult from the
                                  deterministic pass; used to populate
                                  is_sensitive flags and VBA module names.

    Returns:
        WorkbookDigest. On any error, returns a minimal digest with empty
        sheet_summaries and logs the exception.
    """
    file_name = wb_path.name
    vba_present = wb_path.suffix.lower() == ".xlsm"

    try:
        wb = openpyxl.load_workbook(
            str(wb_path),
            data_only=False,
            keep_vba=False,
            read_only=True,
        )

        sensitive_sheets: set[str] = set()
        if workbook_analyser_result is not None:
            sensitive_sheets = set(workbook_analyser_result.high_sensitivity_sheets)

        wb_named_ranges, sheet_named_ranges = _extract_named_ranges(wb)

        sheet_summaries: list[SheetSummary] = []
        for ws in wb.worksheets:
            try:
                summary = _extract_sheet_summary(ws, sensitive_sheets, sheet_named_ranges)
                sheet_summaries.append(summary)
            except Exception as exc:
                logger.warning("digest_builder: failed to summarise sheet '%s': %s", ws.title, exc)

        vba_module_names = _extract_vba_module_names(wb_path, workbook_analyser_result)

        wb.close()

        return WorkbookDigest(
            file_name=file_name,
            sheet_summaries=sheet_summaries,
            workbook_named_ranges=wb_named_ranges,
            vba_present=vba_present,
            vba_module_names=vba_module_names,
        )

    except Exception as exc:
        logger.error("digest_builder: failed to build digest for '%s': %s", file_name, exc)
        return WorkbookDigest(
            file_name=file_name,
            sheet_summaries=[],
            workbook_named_ranges={},
            vba_present=vba_present,
            vba_module_names=[],
        )