"""Kanban search — FTS5 + vector hybrid with process-local cache.

Design constraints (Plan E §F4):
- All kanban fields are Fernet-encrypted at rest; no plaintext on disk.
- The search index lives in a process-local in-memory FTS5 virtual table
  that is rebuilt when the generation counter (bumped on every mutation)
  diverges from the cached generation.
- Vector embeddings (hash-based, from ``utility_model.utility_embed_texts``)
  are stored Fernet-encrypted in a ``kanban_task_vectors`` table on the
  same DB as kanban_tasks. Decrypted vectors are held in the process-local
  cache alongside texts.
- Hybrid results use Reciprocal Rank Fusion (RRF), same pattern as
  ``session_search_db._merge_rows_rrf``.
- Cache is keyed by tenant for multi-tenant isolation.
"""

from __future__ import annotations

import json
import logging
import math
import sqlite3
import threading
from typing import Any

from utils.fts5_query import safe_fts_query

_log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Generation counter — bumped by kanban_storage on every mutation
# ---------------------------------------------------------------------------

_generation_lock = threading.Lock()
_generation: int = 0


def bump_generation() -> None:
    """Called by ``kanban_storage`` after create/update/delete mutations."""
    global _generation
    with _generation_lock:
        _generation += 1


def current_generation() -> int:
    return _generation


# ---------------------------------------------------------------------------
# Per-tenant search cache
# ---------------------------------------------------------------------------

_RRF_K = 60.0
def _kanban_embedding_dims() -> int:
    """Match active embedder dimension (MLX or hash fallback)."""
    try:
        from utils.embedding_model import embedding_dim

        return max(8, int(embedding_dim()))
    except Exception:
        return 256
_MAX_TASK_ROWS = 500

_cache_lock = threading.Lock()
_caches: dict[str | None, "_TenantCache"] = {}


class _TenantCache:
    __slots__ = (
        "gen", "fts_conn", "texts", "vectors", "task_dicts",
    )

    def __init__(self) -> None:
        self.gen: int = -1
        self.fts_conn: sqlite3.Connection | None = None
        self.texts: dict[str, str] = {}
        self.vectors: dict[str, list[float]] = {}
        self.task_dicts: dict[str, dict[str, Any]] = {}

    def close(self) -> None:
        if self.fts_conn:
            try:
                self.fts_conn.close()
            except Exception:
                pass
            self.fts_conn = None


def _get_cache(tenant: str | None) -> _TenantCache:
    with _cache_lock:
        if tenant not in _caches:
            _caches[tenant] = _TenantCache()
        return _caches[tenant]


def invalidate_all_caches() -> None:
    """Force rebuild on next search (used by tests)."""
    with _cache_lock:
        for c in _caches.values():
            c.close()
        _caches.clear()


# ---------------------------------------------------------------------------
# Cache rebuild
# ---------------------------------------------------------------------------

def _rebuild_cache(cache: _TenantCache, tenant: str | None, gen_snapshot: int) -> None:
    """Decrypt all tasks (+ terminal-run summaries) into the FTS5 table and
    vector cache.

    ``gen_snapshot`` is the generation value the caller observed *before*
    starting the rebuild.  We stamp the cache with this value (not a fresh
    ``current_generation()`` read) so that a concurrent ``bump_generation()``
    during the rebuild correctly leaves the cache stale, forcing the next
    caller to rebuild again with the newer data.
    """
    import kanban_storage as ks

    conn = ks._open_conn()
    try:
        where = ["1=1"]
        params: list[Any] = []
        if tenant is not None:
            where.append("tenant=?")
            params.append(tenant)
        sql = (
            f"SELECT * FROM kanban_tasks WHERE {' AND '.join(where)} "
            f"ORDER BY priority ASC, created_at ASC LIMIT {_MAX_TASK_ROWS}"
        )
        rows = conn.execute(sql, params).fetchall()

        new_texts: dict[str, str] = {}
        new_vectors: dict[str, list[float]] = {}
        new_dicts: dict[str, dict[str, Any]] = {}

        for row in rows:
            d = ks._row_to_dict(row)
            if d is None:
                continue
            tid = d.get("id") or ""
            if not tid:
                continue

            title = d.get("title") or ""
            body = d.get("body") or ""
            tags = " ".join(d.get("tags") or [])

            summary = ""
            if d.get("status") in ("done", "cancelled"):
                try:
                    run_row = conn.execute(
                        "SELECT summary_ciphertext FROM kanban_task_runs "
                        "WHERE task_id=? AND outcome IN ('completed','cancelled') "
                        "ORDER BY ended_at DESC LIMIT 1",
                        (tid,),
                    ).fetchone()
                    if run_row and run_row[0]:
                        summary = ks._decrypt(run_row[0])
                except Exception:
                    pass

            corpus = f"{title} {body} {tags} {summary}".strip()
            new_texts[tid] = corpus
            d["parent_ids"] = ks._get_parent_ids(conn, tid)
            new_dicts[tid] = d

        _load_vectors(conn, new_vectors, set(new_texts.keys()))
    finally:
        conn.close()

    fts = sqlite3.connect(":memory:")
    fts.execute("CREATE VIRTUAL TABLE IF NOT EXISTS ktask USING fts5(tid UNINDEXED, content)")
    for tid, corpus in new_texts.items():
        fts.execute("INSERT INTO ktask(tid, content) VALUES (?, ?)", (tid, corpus))

    old_fts = cache.fts_conn
    cache.fts_conn = fts
    cache.texts = new_texts
    cache.vectors = new_vectors
    cache.task_dicts = new_dicts
    cache.gen = gen_snapshot
    if old_fts:
        try:
            old_fts.close()
        except Exception:
            pass


def _load_vectors(
    conn: sqlite3.Connection,
    out: dict[str, list[float]],
    valid_ids: set[str],
) -> None:
    """Load encrypted embeddings from the kanban_task_vectors table."""
    try:
        rows = conn.execute(
            "SELECT task_id, embedding_enc FROM kanban_task_vectors"
        ).fetchall()
    except sqlite3.OperationalError:
        return
    import kanban_storage as ks
    for row in rows:
        tid = row[0]
        if tid not in valid_ids:
            continue
        try:
            raw = ks._decrypt(row[1])
            vec = json.loads(raw)
            expected = _kanban_embedding_dims()
            if isinstance(vec, list) and len(vec) == expected:
                out[tid] = vec
        except Exception:
            continue


# ---------------------------------------------------------------------------
# Ensure vector table schema
# ---------------------------------------------------------------------------

_vector_table_created = False
_vector_table_lock = threading.Lock()


def _ensure_vector_table() -> None:
    global _vector_table_created
    if _vector_table_created:
        return
    with _vector_table_lock:
        if _vector_table_created:
            return
        import kanban_storage as ks
        conn = ks._open_conn()
        try:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS kanban_task_vectors (
                    task_id TEXT PRIMARY KEY,
                    content_hash TEXT NOT NULL DEFAULT '',
                    embedding_enc BLOB NOT NULL,
                    created_at TEXT NOT NULL DEFAULT ''
                )
            """)
            conn.commit()
        finally:
            conn.close()
        _vector_table_created = True


# ---------------------------------------------------------------------------
# Embedding pipeline
# ---------------------------------------------------------------------------

def embed_task(task_id: str) -> bool:
    """Compute and store embedding for a single task. Returns True on success."""
    import kanban_storage as ks

    _ensure_vector_table()
    conn = ks._open_conn()
    try:
        row = conn.execute("SELECT * FROM kanban_tasks WHERE id=?", (task_id,)).fetchone()
        d = ks._row_to_dict(row)
        if d is None:
            return False

        title = d.get("title") or ""
        body = d.get("body") or ""
        tags = " ".join(d.get("tags") or [])

        summary = ""
        if d.get("status") in ("done", "cancelled"):
            try:
                run_row = conn.execute(
                    "SELECT summary_ciphertext FROM kanban_task_runs "
                    "WHERE task_id=? AND outcome IN ('completed','cancelled') "
                    "ORDER BY ended_at DESC LIMIT 1",
                    (task_id,),
                ).fetchone()
                if run_row and run_row[0]:
                    summary = ks._decrypt(run_row[0])
            except Exception:
                pass

        corpus = f"{title} {body} {tags} {summary}".strip()
        if not corpus:
            return False
    finally:
        conn.close()

    try:
        from utility_model import utility_embed_texts
        dims = _kanban_embedding_dims()
        vecs = utility_embed_texts(
            [corpus],
            dimensions=dims,
            activity_operation="kanban_search_embedding",
        )
        if not vecs or not vecs[0]:
            return False
        vec = vecs[0]
        if len(vec) != dims:
            return False
    except Exception:
        _log.debug("utility_embed_texts unavailable for task %s", task_id)
        return False

    import hashlib

    try:
        from utility_model import utility_embedding_identity

        embed_id = utility_embedding_identity()
    except Exception:
        embed_id = ""
    content_hash = hashlib.sha256(
        f"{embed_id}\n{corpus}".encode("utf-8", errors="ignore")
    ).hexdigest()[:16]
    enc = ks._encrypt(json.dumps(vec))
    now = ks._now_iso()

    conn = ks._open_conn()
    try:
        conn.execute(
            "INSERT OR REPLACE INTO kanban_task_vectors "
            "(task_id, content_hash, embedding_enc, created_at) VALUES (?, ?, ?, ?)",
            (task_id, content_hash, enc, now),
        )
        conn.commit()
    finally:
        conn.close()

    bump_generation()
    return True


def embed_pending_tasks() -> int:
    """Embed tasks that have no vector yet. Returns count embedded."""
    import kanban_storage as ks

    _ensure_vector_table()
    conn = ks._open_conn()
    try:
        rows = conn.execute(
            "SELECT kt.id FROM kanban_tasks kt "
            "LEFT JOIN kanban_task_vectors kv ON kt.id = kv.task_id "
            "WHERE kv.task_id IS NULL "
            "LIMIT 50"
        ).fetchall()
    finally:
        conn.close()

    count = 0
    for row in rows:
        if embed_task(row[0]):
            count += 1
    return count


# ---------------------------------------------------------------------------
# FTS search
# ---------------------------------------------------------------------------

def _fts_search(
    cache: _TenantCache,
    query: str,
    limit: int,
) -> list[dict[str, Any]]:
    """Full-text search via the in-memory FTS5 index."""
    if not cache.fts_conn:
        return []

    fts_query = safe_fts_query(query, literal_tokens_only=True)
    if not fts_query:
        return _substring_fallback(cache, query, limit)

    try:
        rows = cache.fts_conn.execute(
            "SELECT tid, rank FROM ktask WHERE ktask MATCH ? ORDER BY rank LIMIT ?",
            (fts_query, limit * 2),
        ).fetchall()
    except sqlite3.OperationalError:
        return _substring_fallback(cache, query, limit)

    results: list[dict[str, Any]] = []
    for row in rows:
        tid = row[0]
        d = cache.task_dicts.get(tid)
        if d:
            out = dict(d)
            out["_search_score"] = abs(row[1])
            out["_source"] = "lexical"
            results.append(out)
    return results[:limit]


def _substring_fallback(
    cache: _TenantCache,
    query: str,
    limit: int,
) -> list[dict[str, Any]]:
    """Fallback substring search when FTS5 query is not viable."""
    q = query.lower()
    results: list[tuple[float, dict[str, Any]]] = []
    for tid, corpus in cache.texts.items():
        if q in corpus.lower():
            d = cache.task_dicts.get(tid)
            if d:
                out = dict(d)
                out["_search_score"] = 1.0
                out["_source"] = "lexical"
                results.append((1.0, out))
    results.sort(key=lambda x: -x[0])
    return [r[1] for r in results[:limit]]


# ---------------------------------------------------------------------------
# Vector search
# ---------------------------------------------------------------------------

def _vector_search(
    cache: _TenantCache,
    query: str,
    limit: int,
) -> list[dict[str, Any]]:
    """Cosine similarity search against cached decrypted vectors."""
    if not cache.vectors:
        return []

    try:
        from utility_model import utility_embed_texts
        dims = _kanban_embedding_dims()
        vecs = utility_embed_texts(
            [query],
            dimensions=dims,
            kind="query",
            activity_operation="kanban_search_query_embedding",
        )
        if not vecs or not vecs[0]:
            return []
        query_vec = vecs[0]
        if len(query_vec) != dims:
            return []
    except Exception:
        return []

    scored: list[tuple[float, str]] = []
    for tid, vec in cache.vectors.items():
        sim = _cosine_similarity(query_vec, vec)
        if sim > 0.01:
            scored.append((sim, tid))
    scored.sort(key=lambda x: -x[0])

    results: list[dict[str, Any]] = []
    for sim, tid in scored[:limit]:
        d = cache.task_dicts.get(tid)
        if d:
            out = dict(d)
            out["_search_score"] = sim
            out["semantic_score"] = sim
            out["_source"] = "semantic"
            results.append(out)
    return results


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    if len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(x * x for x in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)


# ---------------------------------------------------------------------------
# Hybrid RRF merge
# ---------------------------------------------------------------------------

def _merge_rrf(
    fts_results: list[dict[str, Any]],
    vec_results: list[dict[str, Any]],
    limit: int,
) -> list[dict[str, Any]]:
    """Reciprocal Rank Fusion merge of FTS5 and vector results."""
    if not vec_results:
        return fts_results[:limit]
    if not fts_results:
        return vec_results[:limit]

    scored: dict[str, dict[str, Any]] = {}
    for rank, row in enumerate(fts_results, start=1):
        tid = row.get("id", "")
        item = scored.setdefault(tid, {"row": dict(row), "score": 0.0, "sources": set()})
        item["score"] += 1.0 / (_RRF_K + rank)
        item["sources"].add("lexical")

    for rank, row in enumerate(vec_results, start=1):
        tid = row.get("id", "")
        item = scored.setdefault(tid, {"row": dict(row), "score": 0.0, "sources": set()})
        item["score"] += 1.0 / (_RRF_K + rank)
        item["sources"].add("semantic")
        if row.get("semantic_score") and not item["row"].get("semantic_score"):
            item["row"]["semantic_score"] = row["semantic_score"]

    merged = sorted(scored.values(), key=lambda x: -x["score"])
    out: list[dict[str, Any]] = []
    for m in merged[:limit]:
        row = dict(m["row"])
        src = m["sources"]
        if src == {"lexical"}:
            row["_source"] = "lexical"
        elif src == {"semantic"}:
            row["_source"] = "semantic"
        else:
            row["_source"] = "both"
        row["_search_score"] = m["score"]
        out.append(row)
    return out


# ---------------------------------------------------------------------------
# Public search API
# ---------------------------------------------------------------------------

def search(
    query: str,
    *,
    tenant: str | None = None,
    limit: int = 20,
    semantic: bool = False,
) -> list[dict[str, Any]]:
    """Search kanban tasks with FTS5 + optional semantic vectors.

    Returns task dicts with ``_search_score``, ``_source`` ("lexical",
    "semantic", or "both"), and optionally ``semantic_score``.
    """
    q = (query or "").strip()
    if not q:
        return []

    cache = _get_cache(tenant)

    gen = current_generation()
    if cache.gen != gen:
        _rebuild_cache(cache, tenant, gen)

    fts_results = _fts_search(cache, q, limit)

    if semantic and cache.vectors:
        vec_results = _vector_search(cache, q, limit)
        return _merge_rrf(fts_results, vec_results, limit)

    return fts_results


def find_similar(
    task_id: str,
    *,
    tenant: str | None = None,
    limit: int = 5,
) -> list[dict[str, Any]]:
    """Find tasks similar to the given task (by vector similarity)."""
    cache = _get_cache(tenant)

    gen = current_generation()
    if cache.gen != gen:
        _rebuild_cache(cache, tenant, gen)

    vec = cache.vectors.get(task_id)
    if not vec:
        text = cache.texts.get(task_id, "")
        if not text:
            return []
        results = search(text, tenant=tenant, limit=limit + 1, semantic=False)
        return [r for r in results if r.get("id") != task_id][:limit]

    scored: list[tuple[float, str]] = []
    for tid, other_vec in cache.vectors.items():
        if tid == task_id:
            continue
        sim = _cosine_similarity(vec, other_vec)
        if sim > 0.01:
            scored.append((sim, tid))
    scored.sort(key=lambda x: -x[0])

    results: list[dict[str, Any]] = []
    for sim, tid in scored[:limit]:
        d = cache.task_dicts.get(tid)
        if d:
            out = dict(d)
            out["_search_score"] = sim
            out["semantic_score"] = sim
            out["_source"] = "semantic"
            results.append(out)
    return results
