"""Resident E5, Qwen3, PyLate mLateOn and BGE-M3. No per-request model loading."""
import gc
import math
import os
import resource
import threading
from pathlib import Path
from dataclasses import dataclass

from app.embedding import EmbeddingEngine, EmbeddingResult, MODEL_ID as E5_ID, TooManyTokens

QWEN_ID = "Qwen/Qwen3-Embedding-0.6B"
COLBERT_ID = "lightonai/mLateOn"
BGE_M3_ID = "BAAI/bge-m3"
QWEN_MAX = 32768
COLBERT_MAX = 8192
INSTRUCTION = "Given a question, retrieve relevant passages that answer the question"


class ServiceBusy(Exception):
    pass


class ModelUnavailable(Exception):
    """Requested model was intentionally not loaded at startup."""


MODEL_NAMES = ("e5", "qwen", "colbert", "bge-m3")


def enabled_models_from_env():
    """Parse the runtime load allowlist; omitted means all four models."""
    raw = os.getenv("ENABLED_MODELS")
    if raw is None:
        enabled = set(MODEL_NAMES)
        # Retain the earlier BGE-only toggle for existing deployments.
        if os.getenv("BGE_M3_ENABLED", "1") != "1":
            enabled.remove("bge-m3")
        return frozenset(enabled)

    aliases = {
        "e5": "e5", E5_ID: "e5",
        "qwen": "qwen", QWEN_ID: "qwen",
        "colbert": "colbert", "mlateon": "colbert", COLBERT_ID: "colbert",
        "bge-m3": "bge-m3", "bge_m3": "bge-m3", BGE_M3_ID: "bge-m3",
    }
    parts = [part.strip() for part in raw.split(",")]
    if not parts or any(not part for part in parts):
        raise ValueError("ENABLED_MODELS must list at least one model: "
                         + ", ".join(MODEL_NAMES))
    unknown = [part for part in parts if part not in aliases]
    if unknown:
        raise ValueError("Unknown ENABLED_MODELS entries: " + ", ".join(unknown))
    mapped = [aliases[part] for part in parts]
    if len(set(mapped)) != len(mapped):
        raise ValueError("Duplicate models in ENABLED_MODELS")
    return frozenset(mapped)


def memory_mb():
    current = None
    try:
        with open("/proc/self/status", encoding="utf-8") as f:
            for line in f:
                if line.startswith("VmRSS:"):
                    current = round(int(line.split()[1]) / 1024, 1)
                    break
    except OSError:
        pass
    peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return {"rss_mb": current, "peak_rss_mb": round(peak / 1024, 1),
            "budget_mb": int(os.getenv("MEMORY_BUDGET_MB", "15360"))}


def validate_vector(vector, dimensions):
    if len(vector) != dimensions or any(not math.isfinite(v) for v in vector):
        raise RuntimeError("Invalid embedding vector")


class QwenEngine:
    dimensions = 1024
    max_tokens = QWEN_MAX

    def __init__(self, tokenizer, model, torch):
        self.tokenizer, self.model, self.torch = tokenizer, model, torch

    @classmethod
    def load(cls):
        import torch
        from transformers import AutoModel, AutoTokenizer
        path = os.getenv("QWEN_MODEL_DIR", "/opt/models/qwen3-embedding-0.6b")
        tokenizer = AutoTokenizer.from_pretrained(path, padding_side="left",
                                                   local_files_only=True)
        model = AutoModel.from_pretrained(path, torch_dtype=torch.float32,
                                          attn_implementation="sdpa",
                                          local_files_only=True).to("cpu").eval()
        if int(model.config.hidden_size) != 1024:
            raise RuntimeError("Incorrect Qwen output size")
        return cls(tokenizer, model, torch)

    def encode(self, text, kind):
        inp = (f"Instruct: {INSTRUCTION}\nQuery:{text}"
               if kind == "query" else text)
        tokens = self.tokenizer(inp, add_special_tokens=True, truncation=False,
                                return_tensors="pt")
        count = int(tokens["attention_mask"].sum().item())
        if count > self.max_tokens:
            raise TooManyTokens(count, self.max_tokens)
        with self.torch.inference_mode():
            last = self.model(**tokens).last_hidden_state[:, -1, :]
            normalized = self.torch.nn.functional.normalize(last, p=2, dim=1)
            values = normalized[0].float().cpu().tolist()
        validate_vector(values, self.dimensions)
        return EmbeddingResult(tokens_count=count, embedding=values)


@dataclass(frozen=True)
class ColbertResult:
    tokens_count: int
    token_embeddings: list[list[float]]


class ColbertEngine:
    dimensions = 128
    max_tokens = COLBERT_MAX

    def __init__(self, model):
        self.model = model

    @classmethod
    def load(cls):
        from pylate import models
        model = models.ColBERT(
            model_name_or_path=os.getenv("COLBERT_MODEL_DIR", "/opt/models/mlateon"),
            device="cpu", query_length=32, document_length=COLBERT_MAX,
        )
        model.eval()
        model.document_length = COLBERT_MAX
        return cls(model)

    def encode(self, text, kind):
        is_query = kind == "query"
        prefix = self.model.query_prefix if is_query else self.model.document_prefix
        tokens = self.model.tokenizer(f"{prefix}{text}", add_special_tokens=True,
                                      truncation=False)
        count = len(tokens["input_ids"])
        if count > self.max_tokens:
            raise TooManyTokens(count, self.max_tokens)
        if is_query:
            # PyLate pads queries to query_length. Never pad short queries to 8192.
            self.model.query_length = max(32, count)
        output = self.model.encode([text], is_query=is_query, batch_size=1,
                                   show_progress_bar=False)
        if len(output) != 1:
            raise RuntimeError("Unexpected PyLate batch size")
        matrix = output[0]
        if hasattr(matrix, "detach"):
            matrix = matrix.detach().float().cpu()
        if hasattr(matrix, "tolist"):
            matrix = matrix.tolist()
        vectors = [[float(x) for x in row] for row in matrix]
        if not vectors or len(vectors) > COLBERT_MAX:
            raise RuntimeError("Unexpected PyLate output length")
        for row in vectors:
            validate_vector(row, self.dimensions)
        return ColbertResult(tokens_count=len(vectors), token_embeddings=vectors)


class ModelRegistry:
    def __init__(self, e5, qwen, colbert, startup_memory=None, bge_m3=None):
        self.e5, self.qwen, self.colbert = e5, qwen, colbert
        self.bge_m3 = bge_m3
        self.startup_memory = startup_memory or {}
        # Only one CPU inference may execute at a time. Allow a bounded
        # number of other HTTP calls to wait instead of returning 429 instantly.
        # Dify's embeddings provider has a 60-second HTTP timeout.
        pending = max(0, int(os.getenv("INFERENCE_MAX_WAITERS", "16")))
        self._waiting_slots = threading.BoundedSemaphore(pending + 1)
        self._gate = threading.BoundedSemaphore(1)

    @classmethod
    def load(cls):
        import torch
        torch.set_num_threads(max(1, min(4, int(os.getenv("CPU_THREADS", "2")))))
        enabled = enabled_models_from_env()
        readings = {}
        e5 = qwen = colbert = bge_m3 = None
        if "e5" in enabled:
            e5 = EmbeddingEngine.load()
            readings["e5"] = memory_mb()
        if "qwen" in enabled:
            qwen = QwenEngine.load()
            readings["qwen"] = memory_mb()
        if "colbert" in enabled:
            colbert = ColbertEngine.load()
            readings["colbert"] = memory_mb()
        if "bge-m3" in enabled:
            from app.bge_m3 import BgeM3Engine
            bge_m3 = BgeM3Engine.load()
            readings["bge_m3"] = memory_mb()
        gc.collect()
        ram = memory_mb()
        if ram["rss_mb"] is not None and ram["rss_mb"] > ram["budget_mb"]:
            raise MemoryError("All resident models exceed RAM budget")
        return cls(e5, qwen, colbert, readings, bge_m3=bge_m3)

    def loaded_models(self):
        return [name for name, engine in (
            ("e5", self.e5), ("qwen", self.qwen),
            ("colbert", self.colbert), ("bge-m3", self.bge_m3),
        ) if engine is not None]

    def _run(self, callback):
        # Keep inference serialized to protect CPU RAM, but give parallel
        # clients time to queue. The queue is bounded to avoid thread exhaustion.
        if not self._waiting_slots.acquire(blocking=False):
            raise ServiceBusy("Inference queue is full; retry later")
        acquired = False
        try:
            timeout = max(0.0, float(os.getenv("INFERENCE_WAIT_SECONDS", "40")))
            acquired = self._gate.acquire(timeout=timeout)
            if not acquired:
                raise ServiceBusy("Inference queue wait timed out; retry later")
            return callback()
        finally:
            if acquired:
                self._gate.release()
            self._waiting_slots.release()

    def embed(self, text, kind, selected):
        engine = self.qwen if selected in ("qwen", QWEN_ID) else self.e5
        if engine is None:
            raise ModelUnavailable(f"Model {selected or 'e5'} is not enabled")
        result = self._run(lambda: engine.encode(text, kind))
        return (QWEN_ID if engine is self.qwen else E5_ID), engine.dimensions, result

    def openai_encode(self, texts, canonical_model):
        """OpenAI-compatible output is always dense only."""
        from app.batching import batch_encode
        from app.bge_m3 import BGE_M3_ID
        if canonical_model == BGE_M3_ID:
            if self.bge_m3 is None:
                raise ModelUnavailable("Model bge-m3 is not enabled")

            def encode_bge_dense():
                engine = self.bge_m3
                counts = [
                    len(engine.model.tokenizer(text, add_special_tokens=True,
                                               truncation=False)["input_ids"])
                    for text in texts
                ]
                for index, count in enumerate(counts):
                    if count > engine.max_tokens:
                        raise TooManyTokens(count, engine.max_tokens,
                                            input_index=index)
                if sum(counts) > 65536:
                    raise ValueError("Batch exceeds 65536 total input tokens")
                results = []
                group, group_tokens = [], 0
                # Dify may send 16+ texts; internal native BGE max is 8.
                for text, count in zip(texts, counts):
                    if group and (len(group) == 8 or group_tokens + count > 8192):
                        results.extend(engine.encode(
                            group, return_sparse=False, return_colbert=False))
                        group, group_tokens = [], 0
                    group.append(text)
                    group_tokens += count
                if group:
                    results.extend(engine.encode(
                        group, return_sparse=False, return_colbert=False))
                return [EmbeddingResult(
                    tokens_count=item.tokens_count, embedding=item.dense)
                    for item in results]

            return self._run(encode_bge_dense)
        if canonical_model == E5_ID:
            engine, is_qwen = self.e5, False
        elif canonical_model == QWEN_ID:
            engine, is_qwen = self.qwen, True
        else:
            raise ValueError("OpenAI embeddings only support dense E5 and Qwen models")
        if engine is None:
            raise ModelUnavailable(f"Model {canonical_model} is not enabled")
        return self._run(lambda: batch_encode(
            engine, texts, "passage", is_qwen=is_qwen, raw=True,
            truncate_long=(canonical_model == E5_ID),
        ))

    def bge_m3_encode(self, texts, return_colbert=False):
        if self.bge_m3 is None:
            raise ModelUnavailable("Model bge-m3 is not enabled")
        return self._run(lambda: self.bge_m3.encode(
            texts, return_colbert=return_colbert,
        ))

    def colbert_encode(self, text, kind):
        if self.colbert is None:
            raise ModelUnavailable("Model colbert is not enabled")
        return self._run(lambda: self.colbert.encode(text, kind))


    def voyage_encode(self, texts, kind, canonical_model):
        """Atomic inference gate protects all microbatches in one HTTP request."""
        from app.batching import batch_encode
        if canonical_model == E5_ID:
            # Same 512-token safety fallback as the OpenAI E5 endpoint.
            # _dense_batch adds query:/passage: before truncating the tail.
            engine, mode = self.e5, {"truncate_long": True}
        elif canonical_model == QWEN_ID:
            engine, mode = self.qwen, {"is_qwen": True}
        elif canonical_model == COLBERT_ID:
            engine, mode = self.colbert, {"is_colbert": True}
        else:
            raise ValueError("Unknown model")
        if engine is None:
            raise ModelUnavailable(f"Model {canonical_model} is not enabled")
        return self._run(lambda: batch_encode(engine, texts, kind, **mode))