import os
import time
import torch
import numpy as np
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List, Optional, Union
from sentence_transformers import SentenceTransformer

app = FastAPI(title="Qwen3-Embedding-4B API")

model = SentenceTransformer("Qwen/Qwen3-Embedding-4B", device="cpu")
DIM = model.get_sentence_embedding_dimension()


class EmbeddingRequest(BaseModel):
    input: Union[str, List[str]]
    model: str = "Qwen/Qwen3-Embedding-4B"
    encoding_format: Optional[str] = "float"

class EmbeddingObject(BaseModel):
    object: str = "embedding"
    embedding: List[float]
    index: int

class Usage(BaseModel):
    prompt_tokens: int
    total_tokens: int
    duration_ms: float

class EmbeddingResponse(BaseModel):
    object: str = "list"
    data: List[EmbeddingObject]
    model: str
    usage: Usage

@app.post("/v1/embeddings", response_model=EmbeddingResponse)
async def embed(req: EmbeddingRequest):
    texts = req.input if isinstance(req.input, list) else [req.input]
    start = time.time()
    embeddings = model.encode(texts, normalize_embeddings=True, show_progress_bar=False)
    duration = (time.time() - start) * 1000

    data = [
        EmbeddingObject(embedding=e.tolist(), index=i)
        for i, e in enumerate(embeddings)
    ]
    total_tokens = sum(max(1, len(t.split()) * 2) for t in texts)

    return EmbeddingResponse(
        object="list",
        data=data,
        model=req.model,
        usage=Usage(prompt_tokens=total_tokens, total_tokens=total_tokens, duration_ms=round(duration, 2)),
    )

@app.get("/health")
async def health():
    return {"status": "ok", "model": "Qwen3-Embedding-4B", "dim": DIM, "backend": "pytorch-cpu"}