"""
SecureBERT+ embedder — extracted from MurshidUIPipeline.ipynb (cell 15).
Produces a 768-dim float32 embedding for a text paragraph.
Also provides build_text_for_embedding (cell 12).
Original file is NOT modified.
"""

from __future__ import annotations

import numpy as np
from lxml import etree

try:
    import torch
    from transformers import AutoModel, AutoTokenizer
    _TORCH_OK = True
except (ImportError, OSError):
    _TORCH_OK = False

from app.config import settings


def _norm_spaces(s: str) -> str:
    return " ".join((s or "").split()).strip()


def _strip_end_punct(s: str) -> str:
    return (s or "").rstrip(". ").strip()


def build_text_for_embedding(clean_rule: str, summary: str) -> str:
    """Combine LLM summary with rule description — cell 12 of notebook."""
    rule_elem = etree.fromstring(clean_rule.strip())
    raw_desc = rule_elem.findtext("description") or ""
    description = _norm_spaces(raw_desc)
    summary = _norm_spaces(summary)
    description = _norm_spaces(description)

    if not summary and not description:
        return ""
    if summary and not description:
        return summary
    if description and not summary:
        return description

    s0 = _strip_end_punct(summary).lower()
    d0 = _strip_end_punct(description).lower()

    if s0 == d0:
        return _strip_end_punct(summary) + "."
    return f"{_strip_end_punct(summary)}. {_strip_end_punct(description)}."


class SecureBERTEmbedder:
    """Mean-pooling embedder using ehsanaghaei/SecureBERT_Plus — cell 15."""

    MAX_LEN = 512
    BATCH_CHUNKS = 8

    def __init__(self, model_id: str | None = None, device: str | None = None):
        if not _TORCH_OK:
            raise RuntimeError("torch/transformers not available — SecureBERTEmbedder cannot be initialised.")
        mid = model_id or settings.embed_model_id
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        self.tokenizer = AutoTokenizer.from_pretrained(mid, use_fast=True)
        self.model = AutoModel.from_pretrained(mid).to(self.device)
        self.model.eval()
        self.cls_id = self.tokenizer.cls_token_id
        self.sep_id = self.tokenizer.sep_token_id
        self.pad_id = (
            self.tokenizer.pad_token_id
            if self.tokenizer.pad_token_id is not None
            else self.sep_id
        )

    def _chunk_text(self, text: str) -> list[list[int]]:
        token_ids = self.tokenizer.encode(text, add_special_tokens=False)
        chunk_size = self.MAX_LEN - 2
        chunks = []
        for i in range(0, len(token_ids), chunk_size):
            piece = token_ids[i : i + chunk_size]
            chunks.append([self.cls_id] + piece + [self.sep_id])
        return chunks

    def embed_text(self, text: str) -> np.ndarray:
        chunks = self._chunk_text(text)
        all_embs: list[np.ndarray] = []

        for i in range(0, len(chunks), self.BATCH_CHUNKS):
            batch = chunks[i : i + self.BATCH_CHUNKS]
            max_len = max(len(x) for x in batch)
            input_ids, masks = [], []
            for x in batch:
                pad = max_len - len(x)
                input_ids.append(x + [self.pad_id] * pad)
                masks.append([1] * len(x) + [0] * pad)

            ids_t = torch.tensor(input_ids).to(self.device)
            mask_t = torch.tensor(masks).to(self.device)

            with torch.no_grad():
                out = self.model(input_ids=ids_t, attention_mask=mask_t)
                tok_emb = out.last_hidden_state
                mask_exp = mask_t.unsqueeze(-1).expand(tok_emb.size()).float()
                summed = torch.sum(tok_emb * mask_exp, dim=1)
                denom = torch.clamp(mask_exp.sum(dim=1), min=1e-9)
                mean_pooled = summed / denom

            all_embs.append(mean_pooled.cpu().numpy())

        all_embs_np = np.vstack(all_embs)
        para_emb = all_embs_np.mean(axis=0)
        para_emb /= np.linalg.norm(para_emb) + 1e-12
        return para_emb.astype(np.float32)