"""Semantic embedding clustering for block consolidation.

Mean-pool embeddings of a multilingual sentence model build GLOBAL candidate
clusters via cosine blocking + union-find (no chunk loss). Safe pairs (similarity
≥ HARD) are merged without an LLM; borderline pairs in the band [BAND_LOW, HARD) are
handed by the caller to an LLM judge for a yes/no decision. If `transformers`/`torch`
are missing or the model won't load → `embed_sims()` returns `None`, and the caller
falls back to the old panel-judge path (silent deactivation, like the readability gate).

CPU is enough; the caller wraps the blocking inference in `asyncio.to_thread`.
`numpy` comes in transitively via torch (deliberately not in requirements.txt, like torch).
"""

import logging

import numpy as np

from config import EMBEDDING_AKTIV, EMBEDDING_MODELL, EMBEDDING_BLOCK_FLOOR, EMBEDDING_BLOCK_CAP

log = logging.getLogger("creator.embedding")

_model_cache = None       # (tokenizer, model, torch) — singleton
_load_attempt = False     # already tried to load?

EMBEDDING_BATCH = 32      # inference batch size (CPU)
EMBEDDING_MAX_LEN = 128   # title + short description are short → a small truncation cap suffices


def _model():
    """Load the model once. None = clustering off (disabled or load error)."""
    global _model_cache, _load_attempt
    if _load_attempt:
        return _model_cache
    _load_attempt = True
    if not EMBEDDING_AKTIV:
        return None
    try:
        import torch
        from transformers import AutoModel, AutoTokenizer
        tok = AutoTokenizer.from_pretrained(EMBEDDING_MODELL)
        model = AutoModel.from_pretrained(EMBEDDING_MODELL)
        model.eval()
        _model_cache = (tok, model, torch)
        log.info("embedding model loaded: %s", EMBEDDING_MODELL)
    except Exception as e:
        log.warning("embedding clustering disabled (model not loadable): %s", e)
        _model_cache = None
    return _model_cache


def available() -> bool:
    """True if the model could be loaded. Loads on the first call (blocking)."""
    return _model() is not None


def embed(texts: list[str]) -> "np.ndarray | None":
    """Texts → (n, d) L2-normalized, mean-pooled embeddings. None = model off."""
    if _model() is None:
        return None
    tok, model, torch = _model_cache
    out = []
    for i in range(0, len(texts), EMBEDDING_BATCH):
        batch = texts[i:i + EMBEDDING_BATCH]
        enc = tok(batch, return_tensors="pt", truncation=True, max_length=EMBEDDING_MAX_LEN, padding=True)
        with torch.no_grad():
            hidden = model(**enc).last_hidden_state               # (b, t, d)
        mask = enc["attention_mask"].unsqueeze(-1).type_as(hidden)
        vec = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)  # mean-pool without padding
        vec = torch.nn.functional.normalize(vec, p=2, dim=1)       # L2 → cosine = dot product
        out.append(vec.cpu().numpy())
    return np.vstack(out).astype(np.float32)


def _find(parent: list[int], x: int) -> int:
    while parent[x] != x:
        parent[x] = parent[parent[x]]      # Pfad-Kompression
        x = parent[x]
    return x


def _union(parent: list[int], a: int, b: int) -> None:
    ra, rb = _find(parent, a), _find(parent, b)
    if ra != rb:
        parent[max(ra, rb)] = min(ra, rb)  # smallest index = root (deterministic)


def embed_sims(texts: list[str]):
    """Texts → (n, n) cosine matrix · None = model not available (fallback)."""
    embs = embed(texts)
    if embs is None:
        return None
    return embs @ embs.T                    # (n, n) cosine, float32 (~2 MB at n=700)


def capped_blocks(sims, floor: float | None = None, cap: int | None = None) -> list[list[int]]:
    """Coarse similarity blocks for the LLM — high recall, but size-capped.

    Greedy: all pairs with cosine ≥ `floor` in descending cosine order; two blocks are merged
    only if the resulting block stays ≤ `cap`. Prevents the giant component (pure threshold
    blocking would otherwise chain almost everything together) and keeps the LLM lists short.
    → list of blocks (index lists), each node in exactly one block.
    """
    fl = EMBEDDING_BLOCK_FLOOR if floor is None else floor
    cp = EMBEDDING_BLOCK_CAP if cap is None else cap
    n = len(sims)
    parent = list(range(n))
    size = [1] * n
    if n >= 2:
        iu = np.triu_indices(n, k=1)
        s = sims[iu]
        kept = np.where(s >= fl)[0]
        # highest cosine first → the tightest pairs form blocks first
        for k in kept[np.argsort(-s[kept])]:
            i, j = int(iu[0][k]), int(iu[1][k])
            ri, rj = _find(parent, i), _find(parent, j)
            if ri != rj and size[ri] + size[rj] <= cp:
                _union(parent, i, j)
                r = _find(parent, i)
                size[r] = size[ri] + size[rj]
    blocks: dict[int, list[int]] = {}
    for i in range(n):
        blocks.setdefault(_find(parent, i), []).append(i)
    return list(blocks.values())