"""
Column name registry.

Builds a bidirectional alias map per session so that repeated commands
resolve columns in O(1) time. Falls back to rapidfuzz only on first
encounter, then caches the result.
"""
from __future__ import annotations

from typing import Optional
from rapidfuzz import process, fuzz

from config import FUZZY_THRESHOLD


class ColumnRegistry:
    """Session-scoped column name resolver with fuzzy-match cache."""

    def __init__(self) -> None:
        # session_id -> {normalized_alias: real_column_name}
        self._store: dict[str, dict[str, str]] = {}

    # ── public API ──────────────────────────────────────────────────

    def register(self, session_id: str, columns: list[str]) -> None:
        """Build alias map for every column in the dataframe."""
        mapping: dict[str, str] = {}
        for col in columns:
            normalized = col.lower().strip()
            mapping[normalized] = col
            # Common variations the user might type
            mapping[col.replace("_", " ").lower()] = col
            mapping[col.replace(" ", "_").lower()] = col
            mapping[col.replace("-", "_").lower()] = col
            mapping[col.replace("-", " ").lower()] = col
        self._store[session_id] = mapping

    def resolve(self, session_id: str, token: str) -> Optional[str]:
        """Resolve a user-provided token to the real column name.

        1. Exact match on alias map (O(1)).
        2. rapidfuzz WRatio against all aliases.
        3. Tokenize the full command and try sliding-window bigrams.
        """
        mapping = self._store.get(session_id)
        if not mapping:
            return None

        # Direct O(1) lookup
        key = token.lower().strip()
        if key in mapping:
            return mapping[key]

        # Single-token fuzzy
        best = process.extractOne(key, list(mapping.keys()), scorer=fuzz.WRatio)
        if best and best[1] >= FUZZY_THRESHOLD:
            return mapping[best[0]]

        # Sliding-window bigram from the original token (handles multi-word columns)
        parts = token.split()
        for i in range(len(parts) - 1):
            bigram = f"{parts[i]} {parts[i + 1]}".lower()
            if bigram in mapping:
                return mapping[bigram]
            best = process.extractOne(bigram, list(mapping.keys()), scorer=fuzz.WRatio)
            if best and best[1] >= FUZZY_THRESHOLD:
                return mapping[best[0]]

        return None

    def get_columns(self, session_id: str) -> list[str]:
        """Return deduplicated real column names."""
        return list(set(self._store.get(session_id, {}).values()))

    def remove(self, session_id: str) -> None:
        self._store.pop(session_id, None)


# Module-level singleton
column_registry = ColumnRegistry()