"""OpenAI SDK-compatible embeddings API, intentionally separate from Voyage routes.

Configure OpenAI(base_url="https://HOST/openai/v1", api_key="...").
The client is responsible for model-specific query/document prefixes.
"""
from __future__ import annotations

import base64
import logging
import struct
from typing import Literal

from fastapi import APIRouter, Depends
from fastapi.responses import JSONResponse
from pydantic import BaseModel, ConfigDict, Field, field_validator

from app.embedding import MODEL_ID as E5_ID, TooManyTokens
from app.bge_m3 import BGE_M3_ID
from app.models import ModelRegistry, QWEN_ID, ServiceBusy
from app.voyage_api import authenticate, resident_models

router = APIRouter(prefix="/openai/v1", tags=["OpenAI-compatible"])
log = logging.getLogger(__name__)
DenseModelName = Literal[
    "e5", "qwen", "bge-m3", "intfloat/multilingual-e5-small",
    "Qwen/Qwen3-Embedding-0.6B", "BAAI/bge-m3"
]


class OpenAIEmbeddingRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")
    model: DenseModelName
    input: str | list[str] = Field(
        description="Raw input; client must add model-specific query/document prefixes"
    )
    dimensions: int | None = Field(default=None, ge=1)
    encoding_format: Literal["float", "base64"] = "float"
    user: str | None = None
    # Non-standard client probe flag (e.g. OpenAI-compatible integrations).
    # It does not change server-side processing; accept it for compatibility.
    drop_params: bool | None = None

    @field_validator("input")
    @classmethod
    def validate_input(cls, value):
        texts = [value] if isinstance(value, str) else value
        if not texts or len(texts) > 1000:
            raise ValueError("input must contain between 1 and 1000 texts")
        if any(not text.strip() for text in texts):
            raise ValueError("input cannot contain blank strings")
        return value


def api_error(status: int, message: str, error_type="invalid_request_error", code=None):
    return JSONResponse(status_code=status, content={
        "error": {"message": message, "type": error_type, "param": None, "code": code}
    })


@router.post("/embeddings", dependencies=[Depends(authenticate)])
def create_embedding(
    request: OpenAIEmbeddingRequest,
    registry: ModelRegistry = Depends(resident_models),
):
    model = request.model
    if model in ("bge-m3", BGE_M3_ID):
        canonical, engine = BGE_M3_ID, registry.bge_m3
    elif model in ("qwen", QWEN_ID):
        canonical, engine = QWEN_ID, registry.qwen
    else:
        canonical, engine = E5_ID, registry.e5
    if engine is None:
        return api_error(503, f"Model {model} is not enabled", "server_error")
    if request.dimensions is not None and request.dimensions != engine.dimensions:
        log.warning(
            "OpenAI embeddings rejected: reason=unsupported_dimensions model=%s "
            "requested_dimensions=%d expected_dimensions=%d",
            model, request.dimensions, engine.dimensions,
        )
        return api_error(
            400,
            f"Model {model} supports exactly {engine.dimensions} dimensions",
            code="unsupported_dimensions",
        )

    texts = [request.input] if isinstance(request.input, str) else request.input
    try:
        results = registry.openai_encode(texts, canonical)
    except TooManyTokens as exc:
        log.warning(
            "OpenAI embeddings rejected: reason=context_length_exceeded model=%s "
            "batch_size=%d input_index=%s input_tokens=%d max_tokens=%d",
            model, len(texts), exc.input_index, exc.count, exc.limit,
        )
        return api_error(400, str(exc), code="context_length_exceeded")
    except ValueError as exc:
        # Never log input texts or untrusted tokenizer exception contents.
        reason = ("batch_token_budget_exceeded"
                  if str(exc) == "Batch exceeds 65536 total input tokens"
                  else "invalid_batch")
        log.warning(
            "OpenAI embeddings rejected: reason=%s model=%s batch_size=%d",
            reason, model, len(texts),
        )
        return api_error(400, str(exc))
    except ServiceBusy as exc:
        response = api_error(429, str(exc), "rate_limit_error", "rate_limit_exceeded")
        response.headers["Retry-After"] = "2"
        return response
    except Exception:
        log.exception("OpenAI-compatible embedding failed")
        return api_error(500, "Embedding generation failed", "server_error")

    def encode(vector):
        if request.encoding_format == "float":
            return vector
        raw = struct.pack(f"<{len(vector)}f", *vector)
        return base64.b64encode(raw).decode("ascii")

    total_tokens = sum(result.tokens_count for result in results)
    # Compact preview of successful inputs for source attribution.
    previews = [(i, value[:240]) for i, value in enumerate(texts[:3])]
    log.warning("Embeddings input preview: model=%s chunks=%d previews=%r",
                model, len(texts), previews)
    return {
        "object": "list",
        "data": [
            {"object": "embedding", "index": index, "embedding": encode(item.embedding)}
            for index, item in enumerate(results)
        ],
        "model": model,
        "usage": {"prompt_tokens": total_tokens, "total_tokens": total_tokens},
    }


@router.get("/models", dependencies=[Depends(authenticate)])
def list_models(registry: ModelRegistry = Depends(resident_models)):
    available = []
    if registry.e5 is not None:
        available.extend(("e5", E5_ID))
    if registry.qwen is not None:
        available.extend(("qwen", QWEN_ID))
    if registry.bge_m3 is not None:
        available.extend(("bge-m3", BGE_M3_ID))
    return {
        "object": "list",
        "data": [
            {"id": mid, "object": "model", "created": 0, "owned_by": "local"}
            for mid in available
        ],
    }