indexer.py

"""
GCAS Search Engine – Indexer

Responsibilities
----------------
1. Load every .xlsx / .xls file in the configured excel_folder.
2. Serialise each row to a human-readable text string.
3. Embed all rows (local or OpenAI).
4. Build a FAISS IndexFlatIP (cosine-similarity, vectors pre-normalised).
5. Persist indexes + raw data to disk so the server survives restarts without
   having to re-embed on every start.
6. Expose a thread-safe search() helper.
"""
from __future__ import annotations

import logging
import math
import pickle
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

from config import settings
from embeddings import embed_texts

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Module-level in-memory stores
# ---------------------------------------------------------------------------
_index_store: Dict[str, Any] = {}      # table_name -> faiss.Index
_data_store: Dict[str, List[Dict]] = {}  # table_name -> list of cleaned dicts
_file_map: Dict[str, str] = {}          # table_name -> original filename


# ---------------------------------------------------------------------------
# Text serialisation helpers
# ---------------------------------------------------------------------------

# Columns that carry zero semantic value and should be excluded from the
# searchable text string (primary keys, raw numeric IDs, etc.)
_SKIP_COLS = {"PK", "_sheet"}


def _row_to_text(table_name: str, row: Dict[str, Any]) -> str:
    """
    Convert a single DataFrame row to a dense, readable sentence that the
    embedding model can meaningfully encode.

    Example output:
      "Table: CollegeMaster | CollegeName: XYZ College | District: Ahmedabad | ..."
    """
    parts: List[str] = [f"Table: {table_name}"]
    for key, val in row.items():
        if key in _SKIP_COLS:
            continue
        if val is None:
            continue
        # Drop NaN floats
        if isinstance(val, float) and math.isnan(val):
            continue
        s = str(val).strip()
        if s and s.lower() not in ("nan", "none", "nat", ""):
            parts.append(f"{key}: {s}")
    return " | ".join(parts)


def _clean_row(row: Dict[str, Any]) -> Dict[str, Any]:
    """
    Make a row JSON-serialisable:
      - Replace NaN with None
      - Convert numpy scalars to plain Python types
      - Drop internal helper columns
    """
    out: Dict[str, Any] = {}
    for key, val in row.items():
        if key in _SKIP_COLS:
            continue
        if isinstance(val, float) and math.isnan(val):
            out[key] = None
        elif hasattr(val, "item"):          # numpy scalar
            out[key] = val.item()
        elif isinstance(val, (np.bool_,)):
            out[key] = bool(val)
        else:
            out[key] = val
    return out


# ---------------------------------------------------------------------------
# Core indexing
# ---------------------------------------------------------------------------

def load_and_index(excel_folder: Optional[str] = None) -> Dict[str, int]:
    """
    Scan *excel_folder*, load all Excel files, embed every row, and build
    FAISS indexes.  Returns {table_name: row_count} for each indexed table.
    """
    import faiss  # deferred so the module can be imported before faiss is installed

    global _index_store, _data_store, _file_map

    folder = Path(excel_folder or settings.excel_folder)
    if not folder.exists():
        raise FileNotFoundError(f"Excel folder not found: {folder}")

    excel_files = sorted(folder.glob("*.xlsx")) + sorted(folder.glob("*.xls"))
    if not excel_files:
        raise FileNotFoundError(f"No Excel files found in {folder}")

    stats: Dict[str, int] = {}
    new_index_store: Dict[str, Any] = {}
    new_data_store: Dict[str, List[Dict]] = {}
    new_file_map: Dict[str, str] = {}

    for filepath in excel_files:
        table_name = filepath.stem
        logger.info("Indexing '%s' …", filepath.name)
        t0 = time.time()

        try:
            # Load all sheets and concatenate
            xl = pd.ExcelFile(filepath, engine="openpyxl")
            frames: List[pd.DataFrame] = []
            for sheet in xl.sheet_names:
                df = pd.read_excel(xl, sheet_name=sheet)
                df["_sheet"] = sheet
                frames.append(df)
            df = pd.concat(frames, ignore_index=True)

            rows = df.to_dict(orient="records")
            cleaned_rows = [_clean_row(r) for r in rows]
            texts = [_row_to_text(table_name, r) for r in rows]

            logger.info("  Embedding %d rows …", len(texts))
            embeddings = embed_texts(texts)   # (N, D), already L2-normalised

            dim = embeddings.shape[1]
            index = faiss.IndexFlatIP(dim)    # cosine sim (vectors normalised)
            index.add(embeddings)

            new_index_store[table_name] = index
            new_data_store[table_name] = cleaned_rows
            new_file_map[table_name] = filepath.name
            stats[table_name] = len(rows)

            elapsed = time.time() - t0
            logger.info(
                "  ✓ '%s': %d rows indexed in %.1fs", table_name, len(rows), elapsed
            )

        except Exception:
            logger.exception("Failed to index '%s' – skipping.", filepath.name)

    # Atomic swap
    _index_store = new_index_store
    _data_store = new_data_store
    _file_map = new_file_map

    _save_cache()

    # Build entity vocabulary for fuzzy matching / "did you mean?"
    try:
        from fuzzy_matcher import build_college_keyword_map, build_vocabulary
        build_vocabulary(_data_store)
        build_college_keyword_map(_data_store)
    except Exception:
        logger.warning("Could not build fuzzy vocabulary – did you mean? disabled.")

    return stats


# ---------------------------------------------------------------------------
# Disk cache (FAISS index files + pickle for row data)
# ---------------------------------------------------------------------------

def _save_cache() -> None:
    cache_dir = Path(settings.index_cache_folder)
    cache_dir.mkdir(parents=True, exist_ok=True)

    import faiss

    for table_name, index in _index_store.items():
        faiss.write_index(index, str(cache_dir / f"{table_name}.faiss"))

    with open(cache_dir / "data_store.pkl", "wb") as fh:
        pickle.dump(
            {"data": _data_store, "file_map": _file_map},
            fh,
            protocol=pickle.HIGHEST_PROTOCOL,
        )

    logger.info("Index cache saved → %s", cache_dir)


def load_cache() -> bool:
    """
    Try to load FAISS indexes + row data from disk.
    Returns True if successful, False if cache is missing or corrupt.
    """
    global _index_store, _data_store, _file_map

    import faiss

    cache_dir = Path(settings.index_cache_folder)
    pkl_path = cache_dir / "data_store.pkl"

    if not pkl_path.exists():
        logger.info("No index cache found at %s", cache_dir)
        return False

    try:
        with open(pkl_path, "rb") as fh:
            stored = pickle.load(fh)
        _data_store = stored.get("data", {})
        _file_map = stored.get("file_map", {})

        loaded_indexes: Dict[str, Any] = {}
        for table_name in _data_store:
            faiss_path = cache_dir / f"{table_name}.faiss"
            if faiss_path.exists():
                loaded_indexes[table_name] = faiss.read_index(str(faiss_path))
            else:
                logger.warning("Missing FAISS file for table '%s'", table_name)

        _index_store = loaded_indexes
        logger.info(
            "Index loaded from cache: %s (%d total rows)",
            list(_index_store.keys()),
            sum(len(v) for v in _data_store.values()),
        )

        # Rebuild entity vocabulary and college keyword map for fuzzy matching
        try:
            from fuzzy_matcher import build_college_keyword_map, build_vocabulary
            build_vocabulary(_data_store)
            build_college_keyword_map(_data_store)
        except Exception:
            logger.warning("Could not build fuzzy vocabulary.")

        return bool(_index_store)

    except Exception:
        logger.exception("Cache load failed – will reindex.")
        return False


# ---------------------------------------------------------------------------
# Search
# ---------------------------------------------------------------------------

def search(
    query_embedding: np.ndarray,
    top_k: int,
    tables: Optional[List[str]] = None,
) -> List[Dict[str, Any]]:
    """
    Perform FAISS nearest-neighbour search across the requested tables.

    Returns a list of candidate dicts, sorted by descending score:
      {"table": str, "row_index": int, "score": float, "data": dict}
    """
    target_tables = tables if tables else list(_index_store.keys())
    candidates: List[Dict[str, Any]] = []

    qvec = query_embedding.reshape(1, -1).astype(np.float32)

    for table_name in target_tables:
        if table_name not in _index_store:
            logger.warning("Table '%s' not in index – skipping.", table_name)
            continue

        index = _index_store[table_name]
        rows = _data_store[table_name]
        k = min(top_k, len(rows))

        scores, indices = index.search(qvec, k)

        for score, idx in zip(scores[0], indices[0]):
            if idx < 0:  # FAISS returns -1 when the index has fewer than k items
                continue
            candidates.append(
                {
                    "table": table_name,
                    "row_index": int(idx),
                    "score": float(score),
                    "data": rows[int(idx)],
                }
            )

    candidates.sort(key=lambda c: c["score"], reverse=True)
    return candidates


# ---------------------------------------------------------------------------
# Introspection helpers
# ---------------------------------------------------------------------------

def is_ready() -> bool:
    return bool(_index_store)


def get_indexed_tables() -> List[str]:
    return list(_index_store.keys())


def get_total_rows() -> int:
    return sum(len(v) for v in _data_store.values())


def get_schema() -> Dict[str, Dict[str, Any]]:
    schema: Dict[str, Dict[str, Any]] = {}
    for table_name, rows in _data_store.items():
        schema[table_name] = {
            "columns": list(rows[0].keys()) if rows else [],
            "row_count": len(rows),
            "file": _file_map.get(table_name, ""),
        }
    return schema