import logging
import threading
import uuid
from typing import Callable, List

import chromadb
import numpy as np
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings

from back_end.config import (
    EMBEDDING_BATCH_SIZE,
    EMBEDDING_DEVICE,
    EMBEDDING_MAX_SEQ_LENGTH,
    EMBEDDING_MODEL_NAME,
    EMBEDDING_MODEL_REVISION,
    EMBEDDING_ONNX_FILE,
)

logger = logging.getLogger(__name__)

ProgressCallback = Callable[[int, int], None]


class _SFRCodeEmbeddings(Embeddings):
    """Salesforce/SFR-Embedding-Code-400M_R, run with ONNX Runtime.

    Uses the official ONNX export shipped in the model repo (pinned commit), so no
    remote Python code is executed (`trust_remote_code` is not needed), torch is
    not required, and CPU inference is faster. Pooling = CLS token, then L2 norm,
    exactly as the model's sentence-transformers config specifies.
    """

    # instruction prefix specified by the Salesforce AI Research team
    QUERY_INSTRUCTION = "Instruct: Given Code or Text, retrieve relevant content. Query: "

    def __init__(self):
        import onnxruntime as ort
        from huggingface_hub import hf_hub_download
        from tokenizers import Tokenizer

        def fetch(filename: str) -> str:
            return hf_hub_download(EMBEDDING_MODEL_NAME, filename, revision=EMBEDDING_MODEL_REVISION)

        self.tokenizer = Tokenizer.from_file(fetch("tokenizer.json"))
        self.tokenizer.enable_truncation(max_length=EMBEDDING_MAX_SEQ_LENGTH)
        # The shipped tokenizer pads everything to 512; pad to the longest in the batch instead
        self.tokenizer.enable_padding(pad_id=self.tokenizer.token_to_id("[PAD]"), pad_token="[PAD]")

        available = ort.get_available_providers()
        if EMBEDDING_DEVICE != "cpu" and "CUDAExecutionProvider" in available:
            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        else:
            providers = ["CPUExecutionProvider"]

        options = ort.SessionOptions()
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        logger.info("Loading %s (%s) with %s ...", EMBEDDING_MODEL_NAME, EMBEDDING_ONNX_FILE, providers[0])
        self.session = ort.InferenceSession(fetch(EMBEDDING_ONNX_FILE), sess_options=options, providers=providers)
        self.input_names = {i.name for i in self.session.get_inputs()}
        logger.info("Embedding model loaded.")

    def _encode(self, texts: List[str], progress: ProgressCallback | None = None) -> np.ndarray:
        """progress(done_chars, total_chars) is called after every batch and may raise
        to abort. Progress is measured in characters because batches are length-sorted,
        so counting texts would make early (short) batches look unrealistically fast."""
        if not texts:
            return np.zeros((0, 0), dtype=np.float32)
        # Sort by length so each batch pads as little as possible, then restore order
        order = sorted(range(len(texts)), key=lambda i: len(texts[i]))
        result: list[np.ndarray | None] = [None] * len(texts)
        total_chars = sum(len(t) for t in texts) or 1
        done_chars = 0
        if progress:
            progress(0, total_chars)

        for start in range(0, len(texts), EMBEDDING_BATCH_SIZE):
            idx = order[start:start + EMBEDDING_BATCH_SIZE]
            encodings = self.tokenizer.encode_batch([texts[i] for i in idx])
            feeds = {
                "input_ids": np.array([e.ids for e in encodings], dtype=np.int64),
                "attention_mask": np.array([e.attention_mask for e in encodings], dtype=np.int64),
            }
            if "token_type_ids" in self.input_names:
                feeds["token_type_ids"] = np.array([e.type_ids for e in encodings], dtype=np.int64)

            hidden = self.session.run(["last_hidden_state"], feeds)[0]
            cls = hidden[:, 0, :].astype(np.float32)
            cls /= np.clip(np.linalg.norm(cls, axis=1, keepdims=True), 1e-12, None)
            for row, i in enumerate(idx):
                result[i] = cls[row]
            done_chars += sum(len(texts[i]) for i in idx)
            if progress:
                progress(done_chars, total_chars)

        return np.stack(result)

    def embed_documents(self, texts: List[str], progress: ProgressCallback | None = None) -> List[List[float]]:
        return self._encode(texts, progress=progress).tolist()

    def embed_query(self, text: str) -> List[float]:
        # The query MUST have the exact instruction prefix applied before encoding
        return self._encode([self.QUERY_INSTRUCTION + text])[0].tolist()


# Load the model once and share it across sessions (it holds no user data).
# Previously it was reloaded on every repo.
_model_lock = threading.Lock()
_embedding_model: _SFRCodeEmbeddings | None = None


def get_embedding_model() -> _SFRCodeEmbeddings:
    global _embedding_model
    with _model_lock:
        if _embedding_model is None:
            _embedding_model = _SFRCodeEmbeddings()
        return _embedding_model


def build_vector_db(documents: List[Document], progress: ProgressCallback | None = None) -> Chroma:
    """Embed the chunks and store them in a fresh, in-memory vector DB owned by one session.

    In-memory (EphemeralClient) + a random collection name means no shared
    on-disk database: one visitor can never query or overwrite another's data,
    and everything disappears when the session is dropped. `progress` may raise
    to abort (too slow / visitor left) before anything is stored.
    """
    model = get_embedding_model()
    texts = [doc.page_content for doc in documents]
    logger.info("Embedding %d chunks...", len(texts))
    all_embeddings = model.embed_documents(texts, progress=progress)

    client = chromadb.EphemeralClient(
        settings=chromadb.Settings(anonymized_telemetry=False, allow_reset=False)
    )
    vector_db = Chroma(
        client=client,
        collection_name=f"repo_{uuid.uuid4().hex}",
        embedding_function=model,
        collection_metadata={"hnsw:space": "cosine"},
    )
    CHROMA_BATCH_SIZE = 5000
    for i in range(0, len(texts), CHROMA_BATCH_SIZE):
        vector_db._collection.add(
            documents=texts[i:i + CHROMA_BATCH_SIZE],
            metadatas=[doc.metadata for doc in documents[i:i + CHROMA_BATCH_SIZE]],
            embeddings=all_embeddings[i:i + CHROMA_BATCH_SIZE],
            ids=[str(uuid.uuid4()) for _ in documents[i:i + CHROMA_BATCH_SIZE]],
        )
    return vector_db


def drop_vector_db(vector_db: Chroma | None) -> None:
    if vector_db is None:
        return
    try:
        vector_db.delete_collection()
    except Exception:
        pass