tools/spreadsheet.py

"""Deterministic spreadsheet inspection."""

from __future__ import annotations

import json
from pathlib import Path

import pandas as pd
from openpyxl import load_workbook


def describe_workbook(path: str | Path) -> str:
    """Describe sheet dimensions, headers, formulas, and inferred cell data types."""
    workbook = load_workbook(path, data_only=False, read_only=True)
    description: dict[str, object] = {}
    for sheet in workbook.worksheets:
        rows = sheet.iter_rows()
        first_row = next(rows, ())
        headers = [cell.value for cell in first_row]
        formulas = [cell.coordinate for cell in first_row if cell.data_type == "f"]
        types: dict[str, int] = {}
        for cell in first_row:
            types[cell.data_type] = types.get(cell.data_type, 0) + 1
        formula_count = len(formulas)
        for row in rows:  # inspect the complete sheet, not only a preview
            for cell in row:
                types[cell.data_type] = types.get(cell.data_type, 0) + 1
                if cell.data_type == "f":
                    formula_count += 1
                    if len(formulas) < 50:
                        formulas.append(cell.coordinate)
        description[sheet.title] = {
            "rows": sheet.max_row,
            "columns": sheet.max_column,
            "headers": headers,
            "formula_count": formula_count,
            "sample_formula_cells": formulas[:50],
            "cell_types": types,
        }
    return json.dumps(description, ensure_ascii=False, default=str)


def read_sheet(
    path: str | Path, sheet_name: str, max_rows: int = 1000
) -> list[dict[str, object]]:
    frame = pd.read_excel(path, sheet_name=sheet_name)
    return frame.where(pd.notna(frame), None).head(max_rows).to_dict(orient="records")


def filter_rows(
    path: str | Path, sheet_name: str, column: str, value: str, exclude: bool = False
) -> list[dict[str, object]]:
    frame = pd.read_excel(path, sheet_name=sheet_name)
    if column not in frame.columns:
        raise KeyError(f"Unknown column: {column}")
    mask = (
        frame[column].astype(str).str.contains(value, case=False, regex=False, na=False)
    )
    selected = frame[~mask if exclude else mask]
    return selected.where(pd.notna(selected), None).to_dict(orient="records")


def sum_column(
    path: str | Path,
    sheet_name: str,
    column: str,
    filter_column: str | None = None,
    filter_value: str | None = None,
    exclude: bool = False,
) -> float:
    frame = pd.read_excel(path, sheet_name=sheet_name)
    if filter_column and filter_value is not None:
        mask = (
            frame[filter_column]
            .astype(str)
            .str.contains(filter_value, case=False, regex=False, na=False)
        )
        frame = frame[~mask if exclude else mask]
    return float(pd.to_numeric(frame[column], errors="coerce").sum())


def inspect_spreadsheet(path: str | Path, max_rows: int = 200) -> str:
    """Return bounded, structured workbook contents for downstream reasoning."""
    file_path = Path(path)
    if file_path.suffix.lower() not in {".xlsx", ".xls"}:
        raise ValueError("Expected an Excel workbook")
    sheets = pd.read_excel(file_path, sheet_name=None)
    result: dict[str, object] = {
        "description": json.loads(describe_workbook(file_path))
    }
    for name, frame in sheets.items():
        clean = frame.where(pd.notna(frame), None)
        result[str(name)] = {
            "shape": [int(frame.shape[0]), int(frame.shape[1])],
            "columns": [str(column) for column in frame.columns],
            "rows": clean.head(max_rows).to_dict(orient="records"),
            "truncated": len(frame) > max_rows,
        }
    return json.dumps(result, ensure_ascii=False, default=str)


def answer_spreadsheet_question(question: str, path: str | Path) -> str | None:
    """Answer well-defined aggregation questions from workbook data when possible."""
    lowered = question.lower()
    if not (
        "total sales" in lowered
        and ("not including drinks" in lowered or "excluding drinks" in lowered)
    ):
        return None
    describe_workbook(path)  # inspect all sheets, formulas, and cell types first
    frames = pd.read_excel(path, sheet_name=None)
    frame = pd.concat(frames.values(), ignore_index=True)
    normalized = {str(column).strip().lower(): column for column in frame.columns}
    value_column = next(
        (
            normalized[name]
            for name in ("total sales", "sales", "revenue", "amount")
            if name in normalized
        ),
        None,
    )
    if value_column is None:
        beverage = (
            r"\b(?:drink|beverage|soda|cola|coffee|tea|juice|water|shake|"
            r"smoothie|lemonade)\b"
        )
        numeric_columns = [
            column
            for column in frame.columns
            if pd.api.types.is_numeric_dtype(frame[column])
            and not pd.Series([str(column)])
            .str.contains(beverage, case=False, regex=True)
            .iloc[0]
        ]
        if not numeric_columns:
            return None
        total = frame[numeric_columns].apply(pd.to_numeric, errors="coerce").sum().sum()
        return f"${total:,.2f}"
    text_columns = [
        column
        for column in frame.columns
        if not pd.api.types.is_numeric_dtype(frame[column])
    ]
    if not text_columns:
        return None
    beverage = r"\b(?:drink|beverage|soda|cola|coffee|tea|juice|water|shake|smoothie|lemonade)\b"
    drink_mask = pd.Series(False, index=frame.index)
    for column in text_columns:
        drink_mask |= (
            frame[column]
            .astype(str)
            .str.contains(beverage, case=False, regex=True, na=False)
        )
    values = pd.to_numeric(frame[value_column], errors="coerce")
    if values.notna().sum() == 0:
        return None
    return f"${values[~drink_mask].sum():,.2f}"