from __future__ import annotations

import logging
import os
import threading
from typing import Any

import gradio as gr
import spaces

logger = logging.getLogger(__name__)

MODEL_NAME = os.getenv("NLI_MODEL_NAME", "cross-encoder/nli-deberta-v3-base")
NLI_PIPELINE = None
_nli_lock = threading.Lock()


def _get_nli_pipeline():
    global NLI_PIPELINE
    if NLI_PIPELINE is None:
        with _nli_lock:
            if NLI_PIPELINE is not None:
                return NLI_PIPELINE
            try:
                from transformers import pipeline  # type: ignore

                NLI_PIPELINE = pipeline(
                    "text-classification",
                    model=MODEL_NAME,
                    device=-1,
                )
                logger.info("NLI loaded globally: %s", MODEL_NAME)
            except Exception as exc:
                NLI_PIPELINE = None
                logger.warning("NLI unavailable globally (%s).", exc)
                raise
    return NLI_PIPELINE


@spaces.GPU
def score_pairs(pairs: list[dict[str, str]]) -> list[dict[str, Any]]:
    pipeline = _get_nli_pipeline()
    if not isinstance(pairs, list):
        raise ValueError("Expected a list of {'text': premise, 'text_pair': hypothesis} objects.")

    normalized_pairs: list[dict[str, str]] = []
    for pair in pairs:
        if not isinstance(pair, dict):
            raise ValueError("Each pair must be an object.")
        text = pair.get("text")
        text_pair = pair.get("text_pair")
        if not isinstance(text, str) or not isinstance(text_pair, str):
            raise ValueError("Each pair must include string 'text' and 'text_pair' fields.")
        normalized_pairs.append({"text": text, "text_pair": text_pair})

    results = pipeline(normalized_pairs)
    return [
        {"label": str(item.get("label", "")), "score": float(item.get("score", 0.5))}
        for item in results
    ]


demo = gr.Interface(
    fn=score_pairs,
    inputs=gr.JSON(
        label="NLI pairs",
        value=[{"text": "Paris is the capital of France.", "text_pair": "Paris is in France."}],
    ),
    outputs=gr.JSON(label="Pipeline output"),
    api_name="score_pairs",
    title="AFVE NLI Service",
    description="Programmatic NLI scoring service for LLMLens AFVE.",
)


if __name__ == "__main__":
    demo.launch()