import os
import json
import spaces
import gradio as gr
from sentence_transformers import SentenceTransformer

MODEL_ID = os.getenv("LOGOS_EMBED_MODEL", "BAAI/bge-m3")
ENCODE_BATCH_SIZE = max(1, int(os.getenv("LOGOS_EMBED_BATCH_SIZE", "16")))
MAX_INPUTS = max(1, int(os.getenv("LOGOS_EMBED_MAX_INPUTS", "64")))

# Load the model once. ZeroGPU handles the GPU allocation for inference.
MODEL = SentenceTransformer(MODEL_ID, device="cuda")


def _parse_texts(value):
    """Accept either newline-separated text or a JSON list."""
    if isinstance(value, list):
        return [str(x).strip() for x in value if str(x).strip()]

    if value is None:
        return []

    raw = str(value).strip()
    if not raw:
        return []

    # Allow the API/client to send a JSON array while keeping the UI simple.
    if raw.startswith("["):
        try:
            parsed = json.loads(raw)
            if isinstance(parsed, list):
                return [str(x).strip() for x in parsed if str(x).strip()]
        except json.JSONDecodeError:
            pass

    return [line.strip() for line in raw.splitlines() if line.strip()]


@spaces.GPU(duration=120)
def embed_batch(texts):
    cleaned = _parse_texts(texts)

    if len(cleaned) > MAX_INPUTS:
        raise ValueError(f"Too many texts in one request: {len(cleaned)} > {MAX_INPUTS}")

    if not cleaned:
        return {
            "model": MODEL_ID,
            "embeddings": [],
            "dimension": 0,
            "count": 0,
        }

    vectors = MODEL.encode(
        cleaned,
        batch_size=ENCODE_BATCH_SIZE,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )

    return {
        "model": MODEL_ID,
        "embeddings": vectors.tolist(),
        "dimension": int(vectors.shape[1]),
        "count": len(cleaned),
    }


with gr.Blocks(title="LOGOS Embedding Engine") as demo:
    gr.Markdown("# LOGOS Embedding Engine")
    gr.Markdown(
        "Remote BGE-M3 embedding worker for LOGOS. "
        "Enter one text per line, then click Embed."
    )

    texts = gr.Textbox(
        label="Texts",
        lines=10,
        placeholder=(
            "One chunk per line.\n\n"
            "Example:\n"
            "البلوكشين ودوره في توثيق الملكية العقارية\n"
            "رقمنة الإجراءات القضائية وأثرها على العدالة"
        ),
    )

    output = gr.JSON(label="Embedding Result")

    run = gr.Button("Embed", variant="primary")
    run.click(
        embed_batch,
        inputs=texts,
        outputs=output,
        api_name="embed_batch",
    )

demo.launch()