"""Semantic search for GreyMatter.

Uses sqlite-vec for vector similarity and EmbeddingProvider for generating embeddings.
Supports per-table search, cross-table search, and hybrid (FTS5 + vector) search.
"""
import sqlite3
from typing import Any

from .chunking import chunk_text, CHUNK_SIZE
from .db import Database
from .embeddings import EmbeddingProvider, serialize_f32

# Max characters for document-level summary embedding
MAX_DOC_SUMMARY = 2000

# ---------------------------------------------------------------------------
# Mutable registries — extensions register via register_table / register_fts_backend
# ---------------------------------------------------------------------------

# Mutable table config — extensions register via register_table()
_TABLE_CONFIG: dict[str, dict] = {}

# FTS backends for hybrid search — extensions register via register_fts_backend()
_FTS_BACKENDS: dict[str, dict] = {}


def register_table(name: str, config: dict) -> None:
    """Register a table for semantic search indexing.

    Must be called before SemanticSearch is instantiated.

    Args:
        name: Logical table name (e.g., "meetings", "findings")
        config: Dict with keys: vec_table, pk_col, source_table, source_pk,
                text_fn, chunked, and optionally chunk_fn.
    """
    _TABLE_CONFIG[name] = config


def register_fts_backend(
    name: str,
    fts_table: str,
    source_table: str,
    columns: list[str],
) -> None:
    """Register an FTS5 backend for hybrid search.

    Args:
        name: Logical name matching a TABLE_CONFIG key
        fts_table: FTS5 virtual table name (e.g., "meetings_fts")
        source_table: Source table to JOIN against
        columns: FTS5 indexed columns
    """
    _FTS_BACKENDS[name] = {
        "fts_table": fts_table,
        "source_table": source_table,
        "columns": columns,
    }


# ---------------------------------------------------------------------------
# Static table configs — core tables only
# ---------------------------------------------------------------------------

# Core tables — always available
CORE_TABLE_CONFIG = {
    "documents": {
        "vec_table": "vec_documents",
        "pk_col": "document_id",
        "source_table": "documents",
        "source_pk": "id",
        "text_fn": lambda r: " ".join(filter(None, [
            r.get("title", ""),
            (r.get("content") or "")[:MAX_DOC_SUMMARY],
        ])),
        "chunk_fn": lambda r: r.get("content") or "",
        "chunked": True,
    },
    "knowledge": {
        "vec_table": "vec_knowledge",
        "pk_col": "entry_id",
        "source_table": "knowledge_entries",
        "source_pk": "id",
        "text_fn": lambda r: " ".join(filter(None, [
            r.get("title", ""),
            r.get("content", ""),
        ])),
        "chunked": False,
    },
    "observations": {
        "vec_table": "vec_observations",
        "pk_col": "observation_id",
        "source_table": "observations",
        "source_pk": "id",
        "text_fn": lambda r: " ".join(filter(None, [
            r.get("title", ""),
            r.get("content", ""),
        ])),
        "chunked": False,  # observations are short, no chunking needed
    },
    "code_symbols": {
        "vec_table": "vec_code_symbols",
        "pk_col": "symbol_id",
        "source_table": "code_symbols",
        "source_pk": "id",
        "text_fn": lambda r: " ".join(filter(None, [
            r.get("qualified_name", ""),
            r.get("signature", ""),
            r.get("summary", ""),
        ])),
        "chunked": False,  # symbols are short
    },
}

# Initialize the mutable config with core tables only.
# SalesOS tables are registered by the extension via register_table().
_TABLE_CONFIG.update(CORE_TABLE_CONFIG)

# Core FTS backends for hybrid search
CORE_FTS_BACKENDS = {
    "observations": {
        "fts_table": "observations_fts",
        "source_table": "observations",
        "columns": ["title", "content"],
    },
}
_FTS_BACKENDS.update(CORE_FTS_BACKENDS)

# Backward compat alias — external code may import TABLE_CONFIG directly
TABLE_CONFIG = _TABLE_CONFIG



def _vec_upsert(db: "Database", table: str, pk_col: str, row_id: str, vec: bytes) -> None:
    """Upsert a vector into a vec0 virtual table.

    vec0 virtual tables don't support INSERT OR REPLACE, so we use
    DELETE + INSERT to achieve upsert semantics.
    """
    db.execute(f"DELETE FROM {table} WHERE {pk_col} = ?", (row_id,))
    db.execute(
        f"INSERT INTO {table}({pk_col}, embedding) VALUES (?, ?)",
        (row_id, vec),
    )


class SemanticSearch:
    """Vector-based semantic search across GreyMatter data."""

    def __init__(self, db: Database, embedder: EmbeddingProvider):
        self.db = db
        self.embedder = embedder

    def search(
        self,
        query: str,
        tables: list[str] | None = None,
        limit: int = 10,
        detail: str = "chunks",
    ) -> list[dict[str, Any]]:
        """Semantic search across specified tables (or all).

        Args:
            query: Search query text.
            tables: List of table names to search (default: all).
            limit: Maximum number of results.
            detail: 'chunks' returns individual chunk matches with chunk_text;
                    'documents' groups by parent document, keeping best score.
        """
        query_vec = serialize_f32(self.embedder.embed(query))
        search_tables = tables or list(_TABLE_CONFIG.keys())
        all_results: list[dict[str, Any]] = []

        for table_name in search_tables:
            cfg = _TABLE_CONFIG.get(table_name)
            if not cfg:
                continue

            try:
                # --- Layer 1: Document-level search ---
                doc_rows = self.db.execute(
                    f"""SELECT s.*, v.distance
                        FROM {cfg['vec_table']} v
                        JOIN {cfg['source_table']} s ON v.{cfg['pk_col']} = s.{cfg['source_pk']}
                        WHERE v.embedding MATCH ? AND k = ?
                        ORDER BY v.distance""",
                    (query_vec, limit),
                )
                for row in doc_rows:
                    row["source_type"] = table_name
                    row["source_id"] = row[cfg["source_pk"]]
                all_results.extend(doc_rows)
            except sqlite3.OperationalError:
                continue  # source table may not exist (e.g., code_symbols without code package)

            # --- Layer 2: Chunk-level search (for chunked tables) ---
            if cfg.get("chunked"):
                # sqlite-vec MATCH queries don't support JOINs with WHERE on
                # joined columns reliably, so we fetch more chunk results and
                # filter by source_table in Python.
                chunk_rows = self.db.execute(
                    """SELECT v.distance, v.chunk_id
                       FROM vec_chunks v
                       WHERE v.embedding MATCH ? AND k = ?
                       ORDER BY v.distance""",
                    (query_vec, limit * 3),
                )
                if chunk_rows:
                    chunk_ids = [r["chunk_id"] for r in chunk_rows]
                    dist_by_id = {r["chunk_id"]: r["distance"] for r in chunk_rows}

                    # Fetch chunk metadata for matching IDs
                    placeholders = ",".join("?" * len(chunk_ids))
                    chunk_meta = self.db.execute(
                        f"""SELECT id, source_table, source_id, chunk_index, chunk_text
                            FROM embedding_chunks
                            WHERE id IN ({placeholders})
                              AND source_table = ?""",
                        (*chunk_ids, cfg["source_table"]),
                    )

                    # Build chunk results with parent metadata
                    for cm in chunk_meta:
                        parent_row = self.db.execute_one(
                            f"SELECT * FROM {cfg['source_table']} WHERE {cfg['source_pk']} = ?",
                            (cm["source_id"],),
                        )
                        if parent_row:
                            result = dict(parent_row)
                            result["source_type"] = table_name
                            result["chunk_id"] = cm["id"]
                            result["chunk_index"] = cm["chunk_index"]
                            result["chunk_text"] = cm["chunk_text"]
                            result["source_id"] = cm["source_id"]
                            result["distance"] = dist_by_id[cm["id"]]
                            all_results.append(result)

        all_results.sort(key=lambda r: r["distance"])

        if detail == "documents":
            return self._group_by_parent(all_results, limit)

        return all_results[:limit]

    def _group_by_parent(
        self, results: list[dict[str, Any]], limit: int
    ) -> list[dict[str, Any]]:
        """Group results by parent document, keeping best score per parent."""
        seen: dict[str, dict] = {}
        for r in results:
            parent_id = r.get("source_id") or r.get("id", "")
            if parent_id not in seen:
                entry = {k: v for k, v in r.items()
                         if k not in ("chunk_id", "chunk_index")}
                if "chunk_text" in r:
                    entry["best_chunk"] = r["chunk_text"]
                    del entry["chunk_text"]
                seen[parent_id] = entry
            elif "chunk_text" in r and "best_chunk" not in seen[parent_id]:
                # First chunk for a parent whose doc-level row was seen first
                seen[parent_id]["best_chunk"] = r["chunk_text"]
        return list(seen.values())[:limit]

    def index_row(self, table: str, row_id: str, text: str) -> None:
        """Embed and index a single row (document-level).

        For chunked tables, also creates chunk embeddings if text is long.
        Also populates FTS5 index if a backend is registered for this table.
        """
        cfg = _TABLE_CONFIG[table]
        vec = serialize_f32(self.embedder.embed(text))
        _vec_upsert(self.db, cfg["vec_table"], cfg["pk_col"], row_id, vec)

        # Also index into FTS5 if backend is registered
        fts = _FTS_BACKENDS.get(table)
        if fts:
            try:
                col_list = ", ".join(fts["columns"])
                source_row = self.db.execute_one(
                    f"SELECT rowid, {col_list} FROM {fts['source_table']} WHERE {cfg['source_pk']} = ?",
                    (row_id,),
                )
                if source_row:
                    placeholders = ", ".join("?" for _ in fts["columns"])
                    values = [source_row["rowid"]] + [source_row.get(c, "") or "" for c in fts["columns"]]
                    self.db.execute(
                        f"INSERT OR REPLACE INTO {fts['fts_table']}(rowid, {col_list}) VALUES (?, {placeholders})",
                        tuple(values),
                    )
            except Exception:
                pass  # FTS table may not exist yet

        # Chunk if table supports it and text is long
        if cfg.get("chunked") and len(text) > CHUNK_SIZE:
            chunks = chunk_text(text)
            for idx, chunk in enumerate(chunks):
                self.db.execute(
                    "INSERT INTO embedding_chunks (source_table, source_id, chunk_index, chunk_text) VALUES (?, ?, ?, ?)",
                    (cfg["source_table"], row_id, idx, chunk),
                )
                chunk_row = self.db.execute_one("SELECT last_insert_rowid() as id")
                cid = chunk_row["id"]
                cvec = serialize_f32(self.embedder.embed(chunk))
                _vec_upsert(self.db, "vec_chunks", "chunk_id", cid, cvec)

    def index_table(self, table: str) -> int:
        """Batch embed all unindexed rows in a table.

        For tables with chunked=True, also creates chunk-level embeddings.
        """
        cfg = _TABLE_CONFIG[table]

        # --- Document-level embeddings (Layer 1) ---
        rows = self.db.execute(
            f"""SELECT s.* FROM {cfg['source_table']} s
                WHERE s.{cfg['source_pk']} NOT IN (
                    SELECT {cfg['pk_col']} FROM {cfg['vec_table']}
                )"""
        )
        if not rows:
            return 0

        texts = [cfg["text_fn"](r) for r in rows]
        ids = [r[cfg["source_pk"]] for r in rows]

        batch_size = 32
        count = 0
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i : i + batch_size]
            batch_ids = ids[i : i + batch_size]
            vectors = self.embedder.embed_batch(batch_texts)
            for row_id, vec in zip(batch_ids, vectors):
                _vec_upsert(
                    self.db, cfg["vec_table"], cfg["pk_col"],
                    row_id, serialize_f32(vec),
                )
                count += 1

        # --- Chunk-level embeddings (Layer 2) ---
        if cfg.get("chunked"):
            chunk_fn = cfg["chunk_fn"]
            for row in rows:
                row_id = row[cfg["source_pk"]]
                full_text = chunk_fn(row)
                if len(full_text) <= CHUNK_SIZE:
                    continue  # short enough for document-level only

                chunks = chunk_text(full_text)
                if not chunks:
                    continue

                # Store chunks in embedding_chunks
                chunk_ids = []
                for idx, chunk in enumerate(chunks):
                    self.db.execute(
                        "INSERT INTO embedding_chunks (source_table, source_id, chunk_index, chunk_text) VALUES (?, ?, ?, ?)",
                        (cfg["source_table"], row_id, idx, chunk),
                    )
                    chunk_row = self.db.execute_one(
                        "SELECT last_insert_rowid() as id"
                    )
                    chunk_ids.append(chunk_row["id"])

                # Batch embed chunks
                for j in range(0, len(chunks), batch_size):
                    batch_chunks = chunks[j : j + batch_size]
                    batch_chunk_ids = chunk_ids[j : j + batch_size]
                    vectors = self.embedder.embed_batch(batch_chunks)
                    for cid, vec in zip(batch_chunk_ids, vectors):
                        _vec_upsert(
                            self.db, "vec_chunks", "chunk_id",
                            cid, serialize_f32(vec),
                        )

        return count

    def hybrid_search(
        self,
        query: str,
        tables: list[str] | None = None,
        limit: int = 10,
        semantic_weight: float = 0.6,
        detail: str = "chunks",
    ) -> list[dict[str, Any]]:
        """Combine FTS5 keyword + vector semantic results.

        Merges both result sets, normalizes scores, applies weighting.
        Default: 60% semantic, 40% keyword.

        Args:
            detail: 'chunks' returns individual chunk matches;
                    'documents' groups by parent document, keeping best score.
        """
        search_tables = tables or ["meetings"]

        # Vector search — pass detail through so grouping is applied
        vec_results = self.search(query, tables=search_tables, limit=limit * 2, detail=detail)

        # FTS5 keyword search (via registered backends)
        fts_results: list[dict[str, Any]] = []
        for backend_name, backend in _FTS_BACKENDS.items():
            if backend_name in search_tables:
                fts_table = backend["fts_table"]
                source_table = backend["source_table"]
                source_cfg = _TABLE_CONFIG.get(backend_name)
                if not source_cfg:
                    continue
                rows = self.db.execute(
                    f"""SELECT m.*, rank FROM {fts_table} f
                       JOIN {source_table} m ON f.rowid = m.rowid
                       WHERE {fts_table} MATCH ?
                       ORDER BY rank LIMIT ?""",
                    (query, limit * 2),
                )
                for row in rows:
                    row["source_type"] = backend_name
                fts_results.extend(rows)

        # Normalize and merge
        results_by_id: dict[str, dict] = {}

        # Process vector results (distance: lower = better)
        if vec_results:
            max_dist = max(r.get("distance", 0) for r in vec_results) or 1.0
            for r in vec_results:
                dist = r.get("distance", 0)
                norm_score = 1.0 - (dist / max_dist) if max_dist > 0 else 1.0
                rid = r.get("source_id") or r.get("id", r.get("meeting_id", ""))
                results_by_id[rid] = {**r, "vec_score": norm_score, "fts_score": 0.0}

        # Process FTS results (rank: more negative = better match)
        if fts_results:
            min_rank = min(r["rank"] for r in fts_results) or -1.0
            for r in fts_results:
                norm_score = r["rank"] / min_rank if min_rank < 0 else 0.0
                source_type = r.get("source_type", "meetings")
                cfg = _TABLE_CONFIG.get(source_type, {})
                pk = cfg.get("source_pk", "id")
                rid = r.get(pk, r.get("id", ""))
                if rid in results_by_id:
                    results_by_id[rid]["fts_score"] = norm_score
                else:
                    results_by_id[rid] = {
                        **r, "vec_score": 0.0, "fts_score": norm_score,
                        "source_type": source_type,
                    }

        # Compute combined score
        for r in results_by_id.values():
            r["score"] = (
                semantic_weight * r.get("vec_score", 0.0)
                + (1 - semantic_weight) * r.get("fts_score", 0.0)
            )

        # Sort by combined score descending
        ranked = sorted(results_by_id.values(), key=lambda r: r["score"], reverse=True)
        return ranked[:limit]

    def status(self) -> dict[str, Any]:
        """Return embedding counts per table, including chunks."""
        doc_layer = {}
        for name, cfg in _TABLE_CONFIG.items():
            try:
                row = self.db.execute_one(
                    f"SELECT COUNT(*) as c FROM {cfg['vec_table']}"
                )
                doc_layer[name] = row["c"] if row else 0
            except sqlite3.OperationalError:
                doc_layer[name] = 0  # table may not exist

        chunk_row = self.db.execute_one("SELECT COUNT(*) as c FROM vec_chunks")
        chunk_count = chunk_row["c"] if chunk_row else 0

        chunk_breakdown = {}
        for name, cfg in _TABLE_CONFIG.items():
            if cfg.get("chunked"):
                row = self.db.execute_one(
                    "SELECT COUNT(*) as c FROM embedding_chunks WHERE source_table = ?",
                    (cfg["source_table"],),
                )
                chunk_breakdown[name] = row["c"] if row else 0

        return {
            "document_layer": doc_layer,
            "chunk_layer": {
                **chunk_breakdown,
                "total_chunks": chunk_count,
            },
        }
