app/core/layout_extractor.py

"""Raw-grid layout extractor.

Reads the *full cell grid* of a sheet/CSV file WITHOUT assuming the header sits
in row 1. Real-world lab files put headers several rows down, interleave
``mean``/``SEM`` summary rows, and place several tables side-by-side on one
sheet. The classic metadata extractor (``sas_extractor.py``) can't represent
any of that — it only sees row-1 column names.

This module produces, per sheet:
  * a coordinate-addressed grid (``cells[row][col]``), and
  * a compact text rendering with stable row/column indices,

so a local LLM can reason over the literal layout and propose a reshape plan
(which row is the header, which axis is the sample, what to transpose/drop).

Only a capped window of each sheet is rendered to keep the prompt small; the
grids are read on the privacy-safe assumption that the model is self-hosted, so
raw cell values may be shown to it.
"""

import logging
from dataclasses import dataclass
from pathlib import Path

import pandas as pd

logger = logging.getLogger(__name__)

# Caps on what we render to the model. Lab sheets are small; these are generous.
MAX_GRID_ROWS = 60
MAX_GRID_COLS = 40
MAX_CELL_CHARS = 24


@dataclass
class SheetGrid:
    """A single sheet's raw cell grid plus addressing metadata."""

    logical_name: str          # "file.xlsx::Sheet" or "file.csv"
    sheet_name: str | None     # None for non-Excel sources
    num_rows: int              # full used-range row count (before capping)
    num_cols: int              # full used-range col count (before capping)
    cells: list[list[str]]     # cells[row][col], stringified, "" for blanks (capped)
    truncated: bool            # True if the sheet exceeded the render caps


def _stringify(value) -> str:
    """Render a cell value compactly for the model.

    Whole-number floats lose their ``.0`` (858.0 -> "858"), other floats are
    rounded, NaN/None become "", and long strings are truncated.
    """
    if value is None:
        return ""
    if isinstance(value, float):
        if pd.isna(value):
            return ""
        if value.is_integer():
            return str(int(value))
        return str(round(value, 6))
    if isinstance(value, int):
        return str(value)
    text = str(value).strip()
    if text.lower() == "nan":
        return ""
    if len(text) > MAX_CELL_CHARS:
        text = text[: MAX_CELL_CHARS - 1] + "…"
    return text


def _trim_used_range(rows: list[list[str]]) -> tuple[list[list[str]], int, int]:
    """Drop trailing all-empty rows and columns. Returns (grid, num_rows, num_cols)."""
    # Trailing empty rows
    last_row = -1
    for r, row in enumerate(rows):
        if any(cell != "" for cell in row):
            last_row = r
    rows = rows[: last_row + 1]

    if not rows:
        return [], 0, 0

    # Trailing empty columns
    width = max(len(row) for row in rows)
    last_col = -1
    for c in range(width):
        if any(c < len(row) and row[c] != "" for row in rows):
            last_col = c
    num_cols = last_col + 1

    normalized = [[(row[c] if c < len(row) else "") for c in range(num_cols)] for row in rows]
    return normalized, len(normalized), num_cols


def _grid_from_dataframe(df: pd.DataFrame) -> tuple[list[list[str]], int, int, bool]:
    """Convert a header-less DataFrame into a trimmed, capped string grid."""
    rows = [[_stringify(v) for v in row] for row in df.itertuples(index=False, name=None)]
    rows, full_rows, full_cols = _trim_used_range(rows)

    truncated = full_rows > MAX_GRID_ROWS or full_cols > MAX_GRID_COLS
    capped = [row[:MAX_GRID_COLS] for row in rows[:MAX_GRID_ROWS]]
    return capped, full_rows, full_cols, truncated


def extract_grids(file_path: Path, logical_prefix: str | None = None) -> list[SheetGrid]:
    """Extract raw cell grids from a file. One SheetGrid per Excel sheet.

    Args:
        file_path: path to the file (no ``::SheetName`` suffix needed here).
        logical_prefix: display name to use instead of ``file_path.name``
            (e.g. the original upload filename).

    Returns a list of SheetGrid; CSV/SAS yield a single grid, Excel one per sheet.
    """
    name = logical_prefix or file_path.name
    ext = file_path.suffix.lower()

    if ext in (".xlsx", ".xls"):
        return _grids_from_excel(file_path, name, ext)
    if ext == ".csv":
        return [_grid_from_csv(file_path, name)]
    if ext in (".sas7bdat", ".xpt"):
        return [_grid_from_sas(file_path, name)]
    raise ValueError(f"Unsupported file type for layout extraction: {ext}")


def _grids_from_excel(file_path: Path, name: str, ext: str) -> list[SheetGrid]:
    engine = "openpyxl" if ext == ".xlsx" else "xlrd"
    xls = pd.ExcelFile(file_path, engine=engine)
    grids: list[SheetGrid] = []
    for sheet in xls.sheet_names:
        df = pd.read_excel(xls, sheet_name=sheet, header=None, dtype=object)
        cells, full_rows, full_cols, truncated = _grid_from_dataframe(df)
        if full_rows == 0:
            logger.info("Skipping empty sheet '%s' in %s", sheet, name)
            continue
        grids.append(SheetGrid(
            logical_name=f"{name}::{sheet}",
            sheet_name=sheet,
            num_rows=full_rows,
            num_cols=full_cols,
            cells=cells,
            truncated=truncated,
        ))
    if not grids:
        raise ValueError(f"The Excel file '{name}' has no non-empty sheets.")
    return grids


def _grid_from_csv(file_path: Path, name: str) -> SheetGrid:
    from app.core.sas_extractor import _detect_csv_encoding

    encoding = _detect_csv_encoding(file_path)
    try:
        df = pd.read_csv(file_path, header=None, dtype=object, encoding=encoding)
    except (UnicodeDecodeError, LookupError):
        df = pd.read_csv(file_path, header=None, dtype=object, encoding="latin1")
    cells, full_rows, full_cols, truncated = _grid_from_dataframe(df)
    return SheetGrid(
        logical_name=name,
        sheet_name=None,
        num_rows=full_rows,
        num_cols=full_cols,
        cells=cells,
        truncated=truncated,
    )


def _grid_from_sas(file_path: Path, name: str) -> SheetGrid:
    """SAS/XPT files are already clean tables; render header + a few rows as a grid."""
    import pyreadstat

    is_sas7bdat = file_path.suffix.lower() == ".sas7bdat"
    reader = pyreadstat.read_sas7bdat if is_sas7bdat else pyreadstat.read_xport
    try:
        df, _ = reader(str(file_path), row_limit=MAX_GRID_ROWS)
    except UnicodeDecodeError:
        df, _ = reader(str(file_path), row_limit=MAX_GRID_ROWS, encoding="latin1")
    # Prepend the column names as the first grid row so it looks like a sheet.
    header_row = [_stringify(c) for c in df.columns]
    body = [[_stringify(v) for v in row] for row in df.itertuples(index=False, name=None)]
    rows, full_rows, full_cols = _trim_used_range([header_row, *body])
    truncated = full_cols > MAX_GRID_COLS
    capped = [row[:MAX_GRID_COLS] for row in rows[:MAX_GRID_ROWS]]
    return SheetGrid(
        logical_name=name,
        sheet_name=None,
        num_rows=full_rows,
        num_cols=full_cols,
        cells=capped,
        truncated=truncated,
    )


def read_full_grid(physical_path: Path, sheet_name: str | None) -> list[list[str]]:
    """Read a sheet's full cell grid (uncapped) as a 2-D list of strings.

    Coordinates align with the rendered grid (same top-left origin), so a plan's
    row/column indices map directly onto this grid. Used by the reshape executor
    to pull the actual values.
    """
    ext = physical_path.suffix.lower()
    if ext in (".xlsx", ".xls"):
        engine = "openpyxl" if ext == ".xlsx" else "xlrd"
        df = pd.read_excel(
            physical_path, sheet_name=sheet_name or 0, header=None,
            dtype=object, engine=engine,
        )
    elif ext == ".csv":
        from app.core.sas_extractor import _detect_csv_encoding
        encoding = _detect_csv_encoding(physical_path)
        try:
            df = pd.read_csv(physical_path, header=None, dtype=object, encoding=encoding)
        except (UnicodeDecodeError, LookupError):
            df = pd.read_csv(physical_path, header=None, dtype=object, encoding="latin1")
    elif ext in (".sas7bdat", ".xpt"):
        import pyreadstat
        reader = (
            pyreadstat.read_sas7bdat if ext == ".sas7bdat" else pyreadstat.read_xport
        )
        try:
            df, _ = reader(str(physical_path))
        except UnicodeDecodeError:
            df, _ = reader(str(physical_path), encoding="latin1")
        header_row = [_stringify(c) for c in df.columns]
        body = [[_stringify(v) for v in row] for row in df.itertuples(index=False, name=None)]
        return [header_row, *body]
    else:
        raise ValueError(f"Unsupported file type for grid read: {ext}")
    return [[_stringify(v) for v in row] for row in df.itertuples(index=False, name=None)]


def render_grid_text(grid: SheetGrid) -> str:
    """Render a SheetGrid as a coordinate-addressed text block for the LLM prompt.

    Rows are labelled ``r0, r1, …`` and columns ``c0, c1, …`` so the model can
    reference exact coordinates in its reshape plan.
    """
    header = (
        f"### Sheet: {grid.logical_name}  (rows={grid.num_rows}, cols={grid.num_cols}"
        f"{', TRUNCATED' if grid.truncated else ''})"
    )
    if not grid.cells:
        return header + "\n(empty)"

    width = max(len(row) for row in grid.cells)
    col_labels = "      " + " | ".join(f"c{c}" for c in range(width))
    lines = [header, col_labels]
    for r, row in enumerate(grid.cells):
        padded = [(row[c] if c < len(row) else "") for c in range(width)]
        lines.append(f"  r{r:<3} " + " | ".join(padded))
    return "\n".join(lines)