"""Embedding-based semantic theme discovery for knowledge corpora (Phase 6)."""

from __future__ import annotations

import logging
import math
import time
from collections import Counter, defaultdict
from typing import Any

_log = logging.getLogger(__name__)

THEMES_ALGO_VERSION = "v1-kmeans-tfidf"
_KMEANS_MAX_ITER = 25
_MAX_LIST_FIELD = 40


def themes_enabled(cfg: dict[str, Any] | None = None) -> bool:
    """True when semantic theme discovery is enabled (parses string booleans)."""
    section = cfg if cfg is not None else get_themes_config()
    return _cfg_bool(section, "semantic_themes_enabled", True)


def get_themes_config() -> dict[str, Any]:
    """Return [knowledge.themes] with safe defaults."""
    from core.config import get_knowledge_config

    sec = get_knowledge_config().get("themes")
    if not isinstance(sec, dict):
        sec = {}
    out = dict(sec)
    out.setdefault("semantic_themes_enabled", True)
    out.setdefault("min_clusters", 3)
    out.setdefault("max_clusters", 12)
    out.setdefault("min_chunks", 8)
    out.setdefault("sample_cap", 400)
    out.setdefault("label_top_terms", 3)
    out.setdefault("representatives_per_theme", 5)
    out.setdefault("min_coherence", 0.15)
    return out


def _cfg_int(cfg: dict[str, Any], key: str, default: int) -> int:
    raw = cfg.get(key, default)
    try:
        return int(raw)
    except (TypeError, ValueError):
        return default


def _cfg_float(cfg: dict[str, Any], key: str, default: float) -> float:
    raw = cfg.get(key, default)
    try:
        return float(raw)
    except (TypeError, ValueError):
        return default


def _cfg_bool(cfg: dict[str, Any], key: str, default: bool) -> bool:
    raw = cfg.get(key, default)
    if isinstance(raw, bool):
        return raw
    if isinstance(raw, (int, float)):
        return bool(raw)
    if isinstance(raw, str):
        low = raw.strip().lower()
        if low in ("0", "false", "no", "off"):
            return False
        if low in ("1", "true", "yes", "on"):
            return True
    return default


def _tokenizer_refs() -> tuple[Any, set[str]]:
    from tools.knowledge_tools import _TOPIC_FALLBACK_EXTRA_STOPWORDS, _tokenizer_refs as kt_refs

    token_re, stopwords = kt_refs()
    return token_re, set(stopwords) | set(_TOPIC_FALLBACK_EXTRA_STOPWORDS)


def _tokenize_for_themes(text: str) -> list[str]:
    token_re, stopwords = _tokenizer_refs()
    return [t for t in token_re.findall((text or "").lower()) if t and t not in stopwords]


def _fallback_looks_secret(text: str) -> bool:
    lower = (text or "").lower()
    if any(
        marker in lower
        for marker in ("api_key", "apikey", "secret", "password", "authorization:", "bearer ")
    ):
        return True
    return "sk-" in text and len(text) > 24


def _looks_secret(text: str) -> bool:
    try:
        from chat_storage import _looks_secret as fn

        return bool(fn(text))
    except Exception:
        return _fallback_looks_secret(text)


def _expected_embed_identity() -> str:
    from utility_model import utility_embedding_identity

    return utility_embedding_identity()


def _expected_chunker_version() -> str:
    from utils.embedding_chunking import CHUNKER_VERSION

    return CHUNKER_VERSION()


def _graph_vector_dim() -> int:
    try:
        from chat_storage import _vault_embed_dim

        return max(8, int(_vault_embed_dim()))
    except Exception:
        return 256


def _search_vector_dim() -> int:
    import session_search_db as ss

    return max(8, int(ss._vector_dim()))


def _expected_graph_model_id() -> str:
    try:
        from chat_storage import _vault_embed_model_id

        return str(_vault_embed_model_id() or "")
    except Exception:
        return _expected_embed_identity()


def _load_row_source_slug_map(slug: str) -> dict[int, str]:
    """Best-effort row_id -> source_slug from graph table (any freshness)."""
    out: dict[int, str] = {}
    try:
        from chat_storage import _ensure_knowledge_graph_tables, _get_connection

        conn = _get_connection()
        try:
            _ensure_knowledge_graph_tables(conn)
            rows = conn.execute(
                """SELECT row_id, source_slug FROM knowledge_chunk_embeddings
                   WHERE corpus_slug = ?""",
                (slug,),
            ).fetchall()
            for row_id, source_slug in rows:
                try:
                    rid = int(row_id)
                except (TypeError, ValueError):
                    continue
                sid = str(source_slug or "").strip()
                if sid:
                    out[rid] = sid
        finally:
            conn.close()
    except Exception as exc:
        _log.debug("[knowledge_themes] source_slug map failed for %s: %s", slug, exc)
    return out


def _load_fresh_graph_chunks(
    slug: str,
    *,
    embed_identity: str,
    chunker_version: str,
    dim: int,
) -> list[dict[str, Any]]:
    rows_out: list[dict[str, Any]] = []
    try:
        from chat_storage import (
            _decode_vault_embedding_vec,
            _ensure_knowledge_graph_tables,
            _get_connection,
        )

        conn = _get_connection()
        try:
            _ensure_knowledge_graph_tables(conn)
            expected_model = _expected_graph_model_id()
            rows = conn.execute(
                """SELECT row_id, source_slug, label, dim, embedding, model_id,
                          embed_identity, chunker_version
                   FROM knowledge_chunk_embeddings WHERE corpus_slug = ?""",
                (slug,),
            ).fetchall()
            for (
                row_id,
                source_slug,
                label,
                row_dim,
                blob,
                stored_model,
                stored_identity,
                stored_chunker,
            ) in rows:
                if str(stored_identity or "") != embed_identity:
                    continue
                if str(stored_chunker or "") != chunker_version:
                    continue
                if expected_model:
                    stored_mid = str(stored_model or "").strip()
                    if stored_mid and stored_mid != expected_model:
                        continue
                try:
                    rid = int(row_id)
                    row_dim_i = int(row_dim or 0)
                except (TypeError, ValueError):
                    continue
                if rid <= 0:
                    continue
                if row_dim_i != dim:
                    continue
                vec = _decode_vault_embedding_vec(bytes(blob or b""))
                if not vec or len(vec) != dim:
                    continue
                rows_out.append(
                    {
                        "row_id": rid,
                        "vec": vec,
                        "source_slug": str(source_slug or "").strip(),
                        "label": str(label or "").strip(),
                    }
                )
        finally:
            conn.close()
    except Exception as exc:
        _log.debug("[knowledge_themes] graph chunk load failed for %s: %s", slug, exc)
    return rows_out


def _load_search_fallback_chunks(
    slug: str,
    *,
    dim: int,
    source_slug_by_row: dict[int, str],
) -> list[dict[str, Any]]:
    import session_search_db as ss

    thread_id = f"knowledge:{slug}"
    vectors = ss.get_thread_chunk_vectors(thread_id)
    if not vectors:
        return []
    previews = ss.get_thread_chunk_previews(thread_id)
    rows_out: list[dict[str, Any]] = []
    for row_id, (_content_hash, vec) in vectors.items():
        try:
            rid = int(row_id)
        except (TypeError, ValueError):
            continue
        if rid <= 0 or not vec or len(vec) != dim:
            continue
        rows_out.append(
            {
                "row_id": rid,
                "vec": vec,
                "source_slug": source_slug_by_row.get(rid, ""),
                "label": "",
                "preview": previews.get(rid, ""),
            }
        )
    return rows_out


def _downsample_chunks(chunks: list[dict[str, Any]], cap: int) -> list[dict[str, Any]]:
    if len(chunks) <= cap:
        return sorted(chunks, key=lambda c: int(c.get("row_id") or 0))
    ordered = sorted(chunks, key=lambda c: int(c.get("row_id") or 0))
    stride = max(1, len(ordered) // cap)
    sampled = ordered[::stride][:cap]
    if len(sampled) < cap:
        seen = {int(c.get("row_id") or 0) for c in sampled}
        for c in ordered:
            rid = int(c.get("row_id") or 0)
            if rid in seen:
                continue
            sampled.append(c)
            if len(sampled) >= cap:
                break
    return sorted(sampled, key=lambda c: int(c.get("row_id") or 0))


def _kmeans_cosine(
    vectors: Any,
    k: int,
    *,
    rng: Any,
    max_iter: int = _KMEANS_MAX_ITER,
) -> tuple[Any, Any]:
    """Cosine k-means on L2-normalized rows. Returns (labels, centroids)."""
    import numpy as np

    n = int(vectors.shape[0])
    if k <= 0 or n == 0:
        return np.zeros(n, dtype=np.int32), np.zeros((0, vectors.shape[1]), dtype=np.float32)

    k = min(k, n)
    centroids_idx = [int(rng.integers(0, n))]
    for _ in range(1, k):
        centroids = vectors[centroids_idx]
        sims = vectors @ centroids.T
        max_sim = np.max(sims, axis=1)
        dist = 1.0 - max_sim
        dist = np.clip(dist, 0.0, None)
        total = float(dist.sum())
        if total <= 0:
            pick = int(rng.integers(0, n))
        else:
            pick = int(rng.choice(n, p=dist / total))
        if pick not in centroids_idx:
            centroids_idx.append(pick)
        else:
            for candidate in range(n):
                if candidate not in centroids_idx:
                    centroids_idx.append(candidate)
                    break

    centroids = vectors[centroids_idx].copy()
    k_eff = int(centroids.shape[0])
    labels = np.zeros(n, dtype=np.int32)
    for _ in range(max_iter):
        sims = vectors @ centroids.T
        new_labels = np.argmax(sims, axis=1).astype(np.int32)
        if np.array_equal(new_labels, labels):
            break
        labels = new_labels
        for j in range(k_eff):
            mask = labels == j
            if not np.any(mask):
                continue
            mean = vectors[mask].mean(axis=0)
            norm = float(np.linalg.norm(mean))
            centroids[j] = mean / norm if norm > 1e-12 else mean

    return labels, centroids


def _tfidf_top_terms(doc_tokens: list[list[str]], top_n: int) -> list[list[str]]:
    n_docs = len(doc_tokens)
    if n_docs == 0:
        return []
    df: Counter[str] = Counter()
    for tokens in doc_tokens:
        for t in set(tokens):
            df[t] += 1
    out: list[list[str]] = []
    for tokens in doc_tokens:
        tf = Counter(tokens)
        scores: dict[str, float] = {}
        for term, count in tf.items():
            scores[term] = count * math.log((n_docs + 1) / (df[term] + 1))
        ranked = sorted(scores, key=lambda t: (-scores[t], t))
        out.append(ranked[:top_n])
    return out


def _cap_theme_lists(theme: dict[str, Any]) -> dict[str, Any]:
    row = dict(theme)
    row["source_slugs"] = list(row.get("source_slugs") or [])[:_MAX_LIST_FIELD]
    row["representative_row_ids"] = list(row.get("representative_row_ids") or [])[:_MAX_LIST_FIELD]
    row["keywords"] = list(row.get("keywords") or [])[:12]
    return row


def _merge_by_canonical(themes: list[dict[str, Any]]) -> list[dict[str, Any]]:
    merged: dict[str, dict[str, Any]] = {}
    for theme in themes:
        key = str(theme.get("canonical_key") or "").strip()
        if not key:
            continue
        if key not in merged:
            merged[key] = dict(theme)
            continue
        prev = merged[key]
        prev["size"] = int(prev.get("size") or 0) + int(theme.get("size") or 0)
        prev["fraction"] = min(
            1.0,
            float(prev.get("fraction") or 0) + float(theme.get("fraction") or 0),
        )
        if float(theme.get("coherence") or 0) > float(prev.get("coherence") or 0):
            prev["coherence"] = theme.get("coherence")
        slugs = set(prev.get("source_slugs") or []) | set(theme.get("source_slugs") or [])
        prev["source_slugs"] = sorted(slugs)
        reps = list(prev.get("representative_row_ids") or [])
        for rid in theme.get("representative_row_ids") or []:
            if rid not in reps:
                reps.append(rid)
        prev["representative_row_ids"] = reps[:_MAX_LIST_FIELD]
    return list(merged.values())


def discover_semantic_themes(slug: str, *, config: dict[str, Any] | None = None) -> list[dict[str, Any]]:
    """Cluster v2 chunk embeddings; TF-IDF labels. Fail-open -> []."""
    try:
        from utils.knowledge_connections_store import _validate_slug

        safe = _validate_slug(slug)
    except ValueError:
        return []

    cfg = config if config is not None else get_themes_config()
    if not themes_enabled(cfg):
        return []

    t0 = time.monotonic()
    try:
        import numpy as np
    except ImportError:
        _log.warning("[knowledge_themes] numpy unavailable; skipping themes for %s", safe)
        return []

    try:
        out = _discover_semantic_themes_impl(safe, cfg, np=np, t0=t0)
        return out
    except Exception as exc:
        _log.warning("[knowledge_themes] discovery failed for %s: %s", safe, exc)
        return []


def _discover_semantic_themes_impl(
    safe: str,
    cfg: dict[str, Any],
    *,
    np: Any,
    t0: float,
) -> list[dict[str, Any]]:
    min_chunks = max(2, _cfg_int(cfg, "min_chunks", 8))
    min_clusters = max(2, _cfg_int(cfg, "min_clusters", 3))
    max_clusters = max(min_clusters, _cfg_int(cfg, "max_clusters", 12))
    sample_cap = max(min_chunks, _cfg_int(cfg, "sample_cap", 400))
    label_top_terms = max(1, _cfg_int(cfg, "label_top_terms", 3))
    reps_per = max(1, _cfg_int(cfg, "representatives_per_theme", 5))
    min_coherence = max(0.0, _cfg_float(cfg, "min_coherence", 0.15))

    embed_identity = _expected_embed_identity()
    chunker_version = _expected_chunker_version()
    graph_dim = _graph_vector_dim()
    search_dim = _search_vector_dim()

    source_slug_by_row = _load_row_source_slug_map(safe)
    import session_search_db as ss

    previews = ss.get_thread_chunk_previews(f"knowledge:{safe}")

    def _usable_from_raw(raw_chunks: list[dict[str, Any]]) -> list[dict[str, Any]]:
        usable_rows: list[dict[str, Any]] = []
        for ch in raw_chunks:
            try:
                rid = int(ch.get("row_id") or 0)
            except (TypeError, ValueError):
                continue
            if rid <= 0:
                continue
            preview = str(ch.get("preview") or previews.get(rid, ""))
            label = str(ch.get("label") or "")
            doc_text = f"{label}\n{preview}".strip()
            if not doc_text or _looks_secret(doc_text):
                continue
            row = dict(ch)
            row["doc_text"] = doc_text
            usable_rows.append(row)
        return usable_rows

    chunks = _load_fresh_graph_chunks(
        safe,
        embed_identity=embed_identity,
        chunker_version=chunker_version,
        dim=graph_dim,
    )
    usable = _usable_from_raw(chunks)
    if len(usable) < min_chunks:
        search_rows = _load_search_fallback_chunks(
            safe, dim=search_dim, source_slug_by_row=source_slug_by_row
        )
        if search_rows:
            search_usable = _usable_from_raw(search_rows)
            if search_usable:
                if graph_dim == search_dim:
                    by_rid: dict[int, dict[str, Any]] = {}
                    for ch in usable + search_usable:
                        try:
                            rid = int(ch["row_id"])
                        except (TypeError, ValueError, KeyError):
                            continue
                        if rid <= 0:
                            continue
                        by_rid.setdefault(rid, ch)
                    usable = list(by_rid.values())
                else:
                    usable = search_usable

    usable = [
        c
        for c in usable
        if isinstance(c.get("vec"), list) and len(c.get("vec") or []) > 0
    ]
    if len(usable) < min_chunks:
        _log.info(
            "[knowledge_themes] %s: insufficient chunks (%s < %s)",
            safe,
            len(usable),
            min_chunks,
        )
        return []

    sampled = _downsample_chunks(list(usable), sample_cap)
    if not sampled:
        return []
    vec_dims = {len(ch.get("vec") or []) for ch in sampled}
    if len(vec_dims) != 1:
        _log.info(
            "[knowledge_themes] %s: inconsistent embedding dims %s; skipping",
            safe,
            sorted(vec_dims),
        )
        return []
    embed_dim = next(iter(vec_dims))
    if embed_dim <= 0:
        _log.info("[knowledge_themes] %s: zero-length embeddings; skipping", safe)
        return []
    n = len(sampled)
    k = min(max(min_clusters, int(round(math.sqrt(n / 2)))), max_clusters, n)

    try:
        mat = np.array([ch["vec"] for ch in sampled], dtype=np.float32)
    except (KeyError, TypeError, ValueError) as exc:
        _log.info("[knowledge_themes] %s: invalid chunk vectors: %s", safe, exc)
        return []
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    norms = np.where(norms < 1e-12, 1.0, norms)
    mat = mat / norms

    rng = np.random.default_rng(0)
    labels, centroids = _kmeans_cosine(mat, k, rng=rng)

    cluster_indices: dict[int, list[int]] = defaultdict(list)
    for idx, lab in enumerate(labels.tolist()):
        cluster_indices[int(lab)].append(idx)

    doc_tokens_by_cluster: dict[int, list[str]] = {}
    themes_raw: list[dict[str, Any]] = []

    for cluster_id, indices in cluster_indices.items():
        if not indices:
            continue
        cluster_vecs = mat[indices]
        centroid = centroids[cluster_id] if cluster_id < len(centroids) else cluster_vecs.mean(axis=0)
        c_norm = float(np.linalg.norm(centroid))
        if c_norm > 1e-12:
            centroid = centroid / c_norm
        sims = cluster_vecs @ centroid
        coherence = float(np.mean(sims)) if len(sims) else 0.0
        if coherence < min_coherence:
            continue

        order = sorted(
            range(len(indices)),
            key=lambda i: (-float(sims[i]), int(sampled[indices[i]].get("row_id") or 0)),
        )
        rep_row_ids = [
            rid
            for i in order[:reps_per]
            if (rid := int(sampled[indices[i]].get("row_id") or 0)) > 0
        ]
        source_slugs = sorted(
            {
                str(sampled[idx].get("source_slug") or "").strip()
                for idx in indices
                if str(sampled[idx].get("source_slug") or "").strip()
            }
        )
        doc_text = "\n".join(str(sampled[idx].get("doc_text") or "") for idx in indices)
        doc_tokens_by_cluster[cluster_id] = _tokenize_for_themes(doc_text)
        themes_raw.append(
            {
                "cluster_id": cluster_id,
                "size": len(indices),
                "fraction": len(indices) / float(n),
                "coherence": round(coherence, 4),
                "source_slugs": source_slugs,
                "representative_row_ids": rep_row_ids,
                "_indices": indices,
            }
        )

    if not themes_raw:
        return []

    token_lists = [
        doc_tokens_by_cluster.get(int(t["cluster_id"]), []) for t in themes_raw
    ]
    top_terms = _tfidf_top_terms(token_lists, label_top_terms)

    from utils.knowledge_topic_labels import classify_topic_label, normalize_theme_label

    finalized: list[dict[str, Any]] = []
    for theme, keywords in zip(themes_raw, top_terms):
        if not keywords:
            continue
        raw_label = " ".join(keywords)
        if classify_topic_label(raw_label, force_theme=True) != "theme":
            continue
        display, canonical = normalize_theme_label(raw_label)
        if not display or not canonical:
            continue
        finalized.append(
            _cap_theme_lists(
                {
                    "label": display,
                    "canonical_key": canonical,
                    "keywords": keywords,
                    "size": int(theme["size"]),
                    "fraction": round(float(theme["fraction"]), 4),
                    "coherence": float(theme["coherence"]),
                    "source_slugs": theme.get("source_slugs") or [],
                    "representative_row_ids": theme.get("representative_row_ids") or [],
                    "origin": "semantic",
                    "algo_version": THEMES_ALGO_VERSION,
                }
            )
        )

    merged = _merge_by_canonical(finalized)
    merged.sort(key=lambda t: (-int(t.get("size") or 0), str(t.get("canonical_key") or "")))
    out = [_cap_theme_lists(t) for t in merged[:max_clusters]]

    _log.info(
        "[knowledge_themes] %s: chunks=%s k=%s themes=%s elapsed=%.2fs",
        safe,
        n,
        k,
        len(out),
        time.monotonic() - t0,
    )
    return out
