# !/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Author         :  Ailove
------------------------------------
@File           :  app.py.py
@CreateTime     :  2026/7/24 13:54
------------------------------------
"""
import json
import os
import secrets
import tempfile
import time
from typing import Any, Tuple

import gradio as gr
import numpy as np
import spaces
import torch
from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from pydantic import BaseModel, ConfigDict, Field
from transformers import AutoModel, AutoTokenizer

MODEL_NAME = "Qwen/Qwen3-Embedding-8B"
MODEL_ALIASES = {MODEL_NAME, "qwen3-embedding-8b"}
ALLOWED_DIMS = [4096, 2048, 1024, 768, 512, 256, 128]
MAX_DISPLAY_ROWS = 5
MAX_DISPLAY_VALUES = 16
MAX_BATCH_SIZE = int(os.getenv("MAX_BATCH_SIZE", "32"))
MAX_INPUT_CHARS = int(os.getenv("MAX_INPUT_CHARS", "20000"))

TOKENIZER = None
MODEL = None


def _load_model():
    global TOKENIZER, MODEL
    if MODEL is None:
        TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)
        MODEL = AutoModel.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
        MODEL.eval()
        if torch.cuda.is_available():
            MODEL = MODEL.to("cuda")


def _l2_normalize(x: np.ndarray, eps: float = 1e-12) -> np.ndarray:
    norms = np.linalg.norm(x, axis=1, keepdims=True)
    return x / np.maximum(norms, eps)


def _mean_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
    summed = torch.sum(last_hidden_state * mask, dim=1)
    counts = torch.clamp(mask.sum(dim=1), min=1e-9)
    return summed / counts


def _api_tokens() -> list[str]:
    tokens = []
    token_list = os.getenv("EMBEDDINGS_API_TOKENS", "")
    single_token = os.getenv("EMBEDDINGS_API_TOKEN", "")

    tokens.extend(token.strip() for token in token_list.split(",") if token.strip())
    if single_token.strip():
        tokens.append(single_token.strip())

    return tokens


def _token_matches(provided_token: str, expected_tokens: list[str]) -> bool:
    return bool(provided_token) and any(secrets.compare_digest(provided_token, token) for token in expected_tokens)


def _openai_error(
    message: str,
    status_code: int,
    error_type: str = "invalid_request_error",
    param: str | None = None,
    code: str | None = None,
    headers: dict[str, str] | None = None,
) -> JSONResponse:
    return JSONResponse(
        status_code=status_code,
        headers=headers,
        content={
            "error": {
                "message": message,
                "type": error_type,
                "param": param,
                "code": code,
            }
        },
    )


def _raise_openai_error(
    message: str,
    status_code: int,
    error_type: str = "invalid_request_error",
    param: str | None = None,
    code: str | None = None,
    headers: dict[str, str] | None = None,
) -> None:
    raise HTTPException(
        status_code=status_code,
        detail={
            "message": message,
            "type": error_type,
            "param": param,
            "code": code,
        },
        headers=headers,
    )


def verify_bearer_token(authorization: str | None = Header(default=None)) -> None:
    expected_tokens = _api_tokens()
    if not expected_tokens:
        _raise_openai_error(
            "Embedding API token is not configured.",
            status.HTTP_401_UNAUTHORIZED,
            code="invalid_api_key",
            headers={"WWW-Authenticate": "Bearer"},
        )

    prefix = "Bearer "
    if authorization is None or not authorization.startswith(prefix):
        _raise_openai_error(
            "Missing or invalid API key.",
            status.HTTP_401_UNAUTHORIZED,
            code="invalid_api_key",
            headers={"WWW-Authenticate": "Bearer"},
        )

    provided_token = authorization[len(prefix) :].strip()
    if not _token_matches(provided_token, expected_tokens):
        _raise_openai_error(
            "Invalid API key.",
            status.HTTP_401_UNAUTHORIZED,
            code="invalid_api_key",
            headers={"WWW-Authenticate": "Bearer"},
        )


def _parse_ui_texts(text_input: str) -> list[str]:
    texts = [line.strip() for line in (text_input or "").splitlines() if line.strip()]
    if not texts:
        raise gr.Error("Please provide at least one non-empty line of text.")
    return texts


def _validate_texts(texts: list[str]) -> None:
    if not texts:
        _raise_openai_error("Input must contain at least one non-empty string.", status.HTTP_400_BAD_REQUEST, param="input")
    if len(texts) > MAX_BATCH_SIZE:
        _raise_openai_error(
            f"Batch size cannot exceed {MAX_BATCH_SIZE}.",
            status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
            param="input",
            code="request_too_large",
        )
    for index, text in enumerate(texts):
        if not text.strip():
            _raise_openai_error(
                "Input strings must not be empty.",
                status.HTTP_400_BAD_REQUEST,
                param=f"input[{index}]",
            )
        if len(text) > MAX_INPUT_CHARS:
            _raise_openai_error(
                f"Input strings cannot exceed {MAX_INPUT_CHARS} characters.",
                status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
                param=f"input[{index}]",
                code="request_too_large",
            )


@spaces.GPU(duration=300)
def _embed_batch(texts: list[str], dimension: int) -> tuple[list[list[float]], int, float]:
    if dimension not in ALLOWED_DIMS:
        raise ValueError(f"Dimension must be one of {ALLOWED_DIMS}.")

    _load_model()
    start = time.perf_counter()

    with torch.no_grad():
        batch = TOKENIZER(
            texts,
            padding=True,
            truncation=True,
            max_length=8192,
            return_tensors="pt",
        )
        prompt_tokens = int(batch["attention_mask"].sum().item())

        if torch.cuda.is_available():
            batch = {k: v.to("cuda") for k, v in batch.items()}

        outputs = MODEL(**batch)
        pooled = _mean_pool(outputs.last_hidden_state, batch["attention_mask"])
        full_embeddings = pooled.detach().cpu().numpy().astype(np.float32)

    full_embeddings = _l2_normalize(full_embeddings)
    truncated = full_embeddings[:, :dimension]
    truncated = _l2_normalize(truncated)

    elapsed = time.perf_counter() - start
    return truncated.tolist(), prompt_tokens, elapsed


class EmbeddingsRequest(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    model: str = Field(default=MODEL_NAME)
    input: str | list[str]
    dimensions: int = Field(default=4096)
    encoding_format: str = Field(default="float")
    user: str | None = None


def _normalize_api_input(value: str | list[str]) -> list[str]:
    if isinstance(value, str):
        return [value]
    if isinstance(value, list) and all(isinstance(item, str) for item in value):
        return value
    _raise_openai_error("Input must be a string or an array of strings.", status.HTTP_400_BAD_REQUEST, param="input")
    return []


@spaces.GPU(duration=300)
def embed_texts(text_input: str, dimension: int) -> Tuple[str, str, str]:
    if dimension not in ALLOWED_DIMS:
        raise gr.Error(f"Dimension must be one of {ALLOWED_DIMS}.")

    texts = _parse_ui_texts(text_input)
    result_list, _prompt_tokens, elapsed = _embed_batch(texts, dimension)

    display_preview = []
    for row in result_list[:MAX_DISPLAY_ROWS]:
        preview_row = row[:MAX_DISPLAY_VALUES]
        if len(row) > MAX_DISPLAY_VALUES:
            preview_row = preview_row + ["..."]
        display_preview.append(preview_row)

    preview_payload = {
        "preview_embeddings": display_preview,
        "shown_rows": min(len(result_list), MAX_DISPLAY_ROWS),
        "total_rows": len(result_list),
        "shown_values_per_row": MAX_DISPLAY_VALUES,
        "actual_dimension": dimension,
        "note": "Download the JSON file for full embeddings.",
    }

    stats = {
        "model": MODEL_NAME,
        "dimension": dimension,
        "count": len(result_list),
        "inference_time_seconds": round(elapsed, 4),
    }

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as f:
        json.dump(
            {
                "model": MODEL_NAME,
                "dimension": dimension,
                "count": len(result_list),
                "inference_time_seconds": elapsed,
                "embeddings": result_list,
            },
            f,
            ensure_ascii=False,
        )
        file_path = f.name

    return json.dumps(preview_payload, indent=2), json.dumps(stats, indent=2), file_path


api_app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)


@api_app.exception_handler(HTTPException)
async def openai_http_exception_handler(_request: Request, exc: HTTPException) -> JSONResponse:
    if isinstance(exc.detail, dict) and {"message", "type", "param", "code"}.issubset(exc.detail):
        return JSONResponse(
            status_code=exc.status_code,
            headers=exc.headers,
            content={"error": exc.detail},
        )
    return _openai_error(str(exc.detail), exc.status_code)


@api_app.exception_handler(RequestValidationError)
async def openai_validation_exception_handler(_request: Request, exc: RequestValidationError) -> JSONResponse:
    del _request
    first_error = exc.errors()[0] if exc.errors() else {}
    location = first_error.get("loc", [])
    param = ".".join(str(part) for part in location if part != "body") or None
    message = first_error.get("msg", "Invalid request body.")
    return _openai_error(str(message), status.HTTP_400_BAD_REQUEST, param=param, code="invalid_request")


@api_app.post("/v1/embeddings", dependencies=[Depends(verify_bearer_token)])
async def create_embeddings(payload: EmbeddingsRequest) -> dict[str, Any]:
    if payload.model not in MODEL_ALIASES:
        _raise_openai_error(
            f"The model '{payload.model}' does not exist.",
            status.HTTP_404_NOT_FOUND,
            param="model",
            code="model_not_found",
        )
    if payload.dimensions not in ALLOWED_DIMS:
        _raise_openai_error(
            f"Dimensions must be one of {ALLOWED_DIMS}.",
            status.HTTP_400_BAD_REQUEST,
            param="dimensions",
            code="invalid_dimensions",
        )
    if payload.encoding_format != "float":
        _raise_openai_error(
            "Only encoding_format='float' is supported.",
            status.HTTP_400_BAD_REQUEST,
            param="encoding_format",
            code="unsupported_encoding_format",
        )

    texts = _normalize_api_input(payload.input)
    _validate_texts(texts)
    embeddings, prompt_tokens, _elapsed = _embed_batch(texts, payload.dimensions)

    return {
        "object": "list",
        "data": [
            {
                "object": "embedding",
                "index": index,
                "embedding": embedding,
            }
            for index, embedding in enumerate(embeddings)
        ],
        "model": MODEL_NAME,
        "usage": {
            "prompt_tokens": prompt_tokens,
            "total_tokens": prompt_tokens,
        },
    }


DESCRIPTION = """
High-quality text embeddings using **Qwen3-Embedding-8B** — top of MTEB text retrieval benchmarks.
Enter one text per line for batch embedding. L2-normalized, Matryoshka truncation to selected dimension.

Best for: semantic search, RAG, document retrieval, pure text similarity.
"""

FOOTER = (
    "Built by [Xavier Fuentes](https://huggingface.co/xavier-fuentes) @ "
    "[AI Enablement Academy](https://enablement.academy)"
)

with gr.Blocks(title="Text Embeddings - Qwen3-Embedding-8B") as demo:
    gr.Markdown("# Text Embeddings - Qwen3-Embedding-8B")
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        text_input = gr.Textbox(
            label="Input text (one per line for batch)",
            lines=10,
            placeholder="Type one sentence per line...",
        )
        dimension = gr.Dropdown(
            choices=ALLOWED_DIMS,
            value=4096,
            label="Embedding dimension (Matryoshka)",
        )

    with gr.Row():
        embed_btn = gr.Button("Generate Embeddings", variant="primary")
        clear_btn = gr.Button("Clear")

    embeddings_preview = gr.Code(label="Embeddings JSON Preview (truncated)", language="json")
    stats_output = gr.Code(label="Stats", language="json")
    download_file = gr.File(label="Download full embeddings JSON")

    embed_btn.click(
        fn=embed_texts,
        inputs=[text_input, dimension],
        outputs=[embeddings_preview, stats_output, download_file],
        api_name=False,
    )
    clear_btn.click(fn=lambda: ("", "", None), outputs=[embeddings_preview, stats_output, download_file], api_name=False)

    gr.Markdown(FOOTER)

app = gr.mount_gradio_app(api_app, demo.queue(), path="/")