src/queryquest/excel/context.py

"""Excel file discovery and context summarization utilities for prompt building."""

from __future__ import annotations

import hashlib
import json
import re
from pathlib import Path

import pandas as pd

from ..config import EXCEL_DIR

SUPPORTED_SUFFIXES = {".xlsx", ".xls"}


def _table_name_from_file(file_name: str) -> str:
	"""Normalize a workbook or sheet label into a SQL-safe identifier."""
	stem = file_name.rsplit(".", 1)[0]
	normalized = re.sub(r"[^0-9A-Za-z]+", "_", stem).strip("_")
	return normalized or "excel_data"


def _table_name_from_sheet(file_name: str, sheet_name: str) -> str:
	"""Return the SQL table name used by the executor for one workbook sheet."""
	workbook_part = _table_name_from_file(file_name)
	sheet_part = _table_name_from_file(sheet_name)
	return f"{workbook_part}__{sheet_part}"


def normalize_excel_dir(excel_dir: str | Path | None = None) -> Path:
	"""Normalize Excel directory input into an absolute resolved path."""
	path = Path(excel_dir) if excel_dir is not None else EXCEL_DIR
	return path.expanduser().resolve()


def list_excel_files(excel_dir: str | Path | None = None) -> list[Path]:
	"""Return supported Excel workbooks in a deterministic sorted order."""
	directory = normalize_excel_dir(excel_dir)
	if not directory.exists() or not directory.is_dir():
		return []
	return sorted(
		[item for item in directory.iterdir() if item.is_file() and item.suffix.lower() in SUPPORTED_SUFFIXES],
		key=lambda item: item.name.lower(),
	)


def read_excel_file(file_path: Path) -> pd.ExcelFile:
	"""Open an Excel workbook and raise if the expected file is missing."""
	if not file_path.exists():
		raise FileNotFoundError(f"File {file_path.name} not found in {file_path.parent}")
	return pd.ExcelFile(file_path)


def get_col_info(df: pd.DataFrame) -> str:
	"""Build compact schema metadata for a DataFrame."""
	col_types = df.dtypes
	rows = len(df)
	cols = len(df.columns)
	info = "\n".join([f"{col}: {dtype}" for col, dtype in col_types.items()])
	if not info:
		info = "No columns found"
	info += f"\nColumns: {cols}"
	info += f"\nRows: {rows}"
	return info


def get_sample_rows(df: pd.DataFrame, sample_size: int = 1) -> str:
	"""Return a small human-readable table preview used in LLM context text."""
	if df.empty:
		return "Sample rows: No rows found"

	sample_df = df.head(sample_size)
	headers = [str(column).strip() for column in sample_df.columns]
	lines = [" | ".join(headers)]

	for row in sample_df.itertuples(index=False, name=None):
		values = [str(value).strip() for value in row]
		lines.append(" | ".join(values))

	return "Sample rows:\n" + "\n".join(lines)


def _snapshot_file(file_path: Path) -> dict[str, str | int]:
	"""Return file metadata used to detect workbook set/content changes."""
	stat = file_path.stat()
	return {
		"name": file_path.name,
		"size": stat.st_size,
		"mtime_ns": stat.st_mtime_ns,
	}


def get_excel_snapshot(excel_dir: str | Path | None = None) -> tuple[str, list[dict[str, str | int]]]:
	"""Return snapshot digest and entries for all tracked Excel files."""
	files = list_excel_files(excel_dir)
	snapshot = [_snapshot_file(file_path) for file_path in files]
	digest = hashlib.sha256(json.dumps(snapshot, sort_keys=True).encode("utf-8")).hexdigest()
	return digest, snapshot


def describe_excel_snapshot_changes(
	previous_snapshot: list[dict[str, str | int]] | None,
	current_snapshot: list[dict[str, str | int]] | None,
) -> str:
	"""Describe added/removed workbook names between two snapshots."""
	previous_names = {
		str(item.get("name"))
		for item in (previous_snapshot or [])
		if isinstance(item, dict) and isinstance(item.get("name"), str)
	}
	current_names = {
		str(item.get("name"))
		for item in (current_snapshot or [])
		if isinstance(item, dict) and isinstance(item.get("name"), str)
	}

	added = sorted(current_names - previous_names)
	removed = sorted(previous_names - current_names)
	if not added and not removed:
		return ""

	sections: list[str] = ["### EXCEL FILE CHANGES"]
	if added:
		sections.append(f"Added files: {', '.join(added)}")
	if removed:
		sections.append(f"Removed files: {', '.join(removed)}")
	return "\n".join(sections)


def build_excel_files_info(excel_dir: str | Path | None = None) -> tuple[str, str, list[dict[str, str | int]]]:
	"""Build text context, snapshot digest, and snapshot entries for prompts."""
	directory = normalize_excel_dir(excel_dir)
	files = list_excel_files(directory)
	digest, snapshot = get_excel_snapshot(directory)

	if not files:
		message = f"Excel directory: {directory}\nNo Excel files found."
		return message, digest, snapshot

	sections: list[str] = [f"Excel directory: {directory}", f"Excel files found: {len(files)}"]
	for file_path in files:
		xls = read_excel_file(file_path)
		sheet_names = [str(sheet_name) for sheet_name in xls.sheet_names]
		sheet_sections: list[str] = [f"File: {file_path.name}", f"Sheets: {', '.join(sheet_names)}"]
		for sheet_name in xls.sheet_names:
			df = xls.parse(sheet_name)
			sheet_sections.append(f"Sheet: {sheet_name}")
			sheet_sections.append(f"SQL table: {_table_name_from_sheet(file_path.name, str(sheet_name))}")
			sheet_sections.append(get_col_info(df))
			sheet_sections.append(get_sample_rows(df))
		sections.append("\n".join(sheet_sections))

	return "\n\n".join(sections), digest, snapshot


def format_excel_context(info: str | None) -> str:
	"""Wrap plain Excel metadata text in a structured prompt section."""
	if not info:
		return ""
	return f"### EXCEL FILE CONTEXT\n{info}"