from __future__ import annotations

from typing import Iterator, List

import duckdb
import numpy as np
import pandas as pd

from blazematch.config import BlockingRule, EmbeddingBlockingConfig


class RuleBlocker:
    """SQL-based blocking using DuckDB for efficient in-process joins.

    Warning:
        Blocking rules are raw SQL expressions interpolated into queries.
        Only use trusted rule strings — never pass unsanitized user input
        as a BlockingRule.
    """

    # Tokens that should never appear in a blocking rule expression
    _BLOCKED_TOKENS = frozenset({
        "drop", "delete", "insert", "update", "alter", "create",
        "truncate", "exec", "execute", "grant", "revoke", "--", "/*",
    })

    def __init__(
        self,
        df_left: pd.DataFrame,
        df_right: pd.DataFrame,
        rules: List[BlockingRule],
        max_candidates_per_record: int = 100,
    ):
        self._df_left = df_left.reset_index(drop=True)
        self._df_right = df_right.reset_index(drop=True)
        self._rules = rules
        self._max_candidates = max_candidates_per_record

        # Validate rules at construction time
        for rule in rules:
            self._validate_rule(rule.sql)

    @classmethod
    def _validate_rule(cls, sql: str) -> None:
        """Reject rules containing dangerous SQL tokens."""
        lowered = sql.lower()
        for token in cls._BLOCKED_TOKENS:
            if token in lowered:
                raise ValueError(
                    f"Blocking rule contains disallowed token '{token}': {sql!r}. "
                    "Rules must be simple join conditions like 'l.col = r.col'."
                )
        # Rules must reference both l and r tables
        if "l." not in lowered and "r." not in lowered:
            raise ValueError(
                f"Blocking rule must reference tables 'l' and 'r': {sql!r}. "
                "Example: 'l.postcode = r.postcode'."
            )

    def block(self) -> pd.DataFrame:
        """Generate candidate pairs by executing all blocking rules.

        Returns a DataFrame with columns [left_idx, right_idx].
        """
        df_l = self._df_left.copy()
        df_l["__left_idx"] = range(len(df_l))
        df_r = self._df_right.copy()
        df_r["__right_idx"] = range(len(df_r))

        con = duckdb.connect()
        try:
            con.register("l", df_l)
            con.register("r", df_r)

            queries = []
            for rule in self._rules:
                queries.append(
                    f"SELECT DISTINCT l.__left_idx AS left_idx, "
                    f"r.__right_idx AS right_idx "
                    f"FROM l JOIN r ON {rule.sql}"
                )

            union_query = " UNION ".join(queries)

            final_query = f"""
                WITH all_pairs AS ({union_query}),
                ranked AS (
                    SELECT left_idx, right_idx,
                        ROW_NUMBER() OVER (PARTITION BY left_idx ORDER BY right_idx) AS rn
                    FROM all_pairs
                )
                SELECT left_idx, right_idx FROM ranked
                WHERE rn <= {self._max_candidates}
            """

            result = con.execute(final_query).fetchdf()
        finally:
            con.close()
        return result

    def block_iter(self, chunk_size: int = 2_000_000) -> Iterator[pd.DataFrame]:
        """Yield candidate pairs in chunks for large-scale processing."""
        pairs = self.block()
        for start in range(0, len(pairs), chunk_size):
            yield pairs.iloc[start : start + chunk_size].reset_index(drop=True)


class EmbeddingBlocker:
    """ANN-based blocking using TF-IDF character n-grams + FAISS.

    Generates character n-gram TF-IDF vectors for specified text fields,
    builds a FAISS index on the right dataset, and retrieves top-k
    nearest neighbors for each left record.

    Requires: pip install blazematch[ann]  (installs faiss-cpu)
    """

    def __init__(
        self,
        df_left: pd.DataFrame,
        df_right: pd.DataFrame,
        config: EmbeddingBlockingConfig,
        max_candidates_per_record: int = 100,
    ):
        self._df_left = df_left.reset_index(drop=True)
        self._df_right = df_right.reset_index(drop=True)
        self._config = config
        self._max_candidates = max_candidates_per_record

    def _build_text(self, df: pd.DataFrame) -> List[str]:
        """Combine configured fields into a single text per record."""
        if not self._config.fields:
            raise ValueError(
                "EmbeddingBlockingConfig.fields must not be empty."
            )
        parts = []
        for field in self._config.fields:
            parts.append(df[field].astype(str).fillna(""))
        return [" ".join(vals) for vals in zip(*parts)]

    def block(self) -> pd.DataFrame:
        """Generate candidate pairs using ANN search.

        Returns a DataFrame with columns [left_idx, right_idx].
        """
        try:
            import faiss
        except ImportError:
            raise ImportError(
                "faiss-cpu is required for EmbeddingBlocker. "
                "Install with: pip install blazematch[ann]"
            )

        from sklearn.feature_extraction.text import TfidfVectorizer

        # Build TF-IDF vectors
        text_left = self._build_text(self._df_left)
        text_right = self._build_text(self._df_right)

        vectorizer = TfidfVectorizer(
            analyzer="char_wb",
            ngram_range=self._config.ngram_range,
            max_features=self._config.max_features,
        )

        # Fit on right side, transform both — avoids concatenating lists
        all_text = text_left + text_right
        vectorizer.fit(all_text)

        sparse_left = vectorizer.transform(text_left)
        sparse_right = vectorizer.transform(text_right)

        # Convert right side to dense for FAISS index building, then free sparse
        vecs_right = np.ascontiguousarray(sparse_right.toarray(), dtype=np.float32)
        del sparse_right
        faiss.normalize_L2(vecs_right)

        # Build FAISS index on right side, then free the dense copy
        dim = vecs_right.shape[1]
        index = faiss.IndexFlatIP(dim)  # inner product (= cosine on normalized)
        index.add(vecs_right)
        del vecs_right

        # Search left side in chunks to limit peak memory:
        # only one chunk's dense matrix is in memory at a time.
        top_k = min(self._config.top_k, len(self._df_right))
        n_left = sparse_left.shape[0]
        chunk_size = max(1, min(10_000, n_left))
        dist_chunks = []
        idx_chunks = []

        for start in range(0, n_left, chunk_size):
            end = min(start + chunk_size, n_left)
            chunk_dense = np.ascontiguousarray(
                sparse_left[start:end].toarray(), dtype=np.float32
            )
            faiss.normalize_L2(chunk_dense)
            D, I = index.search(chunk_dense, top_k)
            dist_chunks.append(D)
            idx_chunks.append(I)
            del chunk_dense

        del sparse_left
        distances = np.vstack(dist_chunks)
        indices = np.vstack(idx_chunks)
        del dist_chunks, idx_chunks

        # Vectorized pair building (avoids Python for-loop over all pairs)
        left_grid = np.repeat(np.arange(n_left, dtype=np.int64), top_k)
        right_flat = indices.ravel()
        sim_flat = distances.ravel()
        mask = (right_flat >= 0) & (sim_flat >= self._config.min_sim)

        pairs = pd.DataFrame({
            "left_idx": left_grid[mask],
            "right_idx": right_flat[mask],
        })

        # Deduplicate and cap
        pairs = pairs.drop_duplicates()
        if self._max_candidates < len(self._df_right):
            pairs = (
                pairs.groupby("left_idx")
                .head(self._max_candidates)
                .reset_index(drop=True)
            )

        return pairs

    def block_iter(self, chunk_size: int = 2_000_000) -> Iterator[pd.DataFrame]:
        """Yield candidate pairs in chunks."""
        pairs = self.block()
        for start in range(0, len(pairs), chunk_size):
            yield pairs.iloc[start : start + chunk_size].reset_index(drop=True)
