import os
import time
from typing import List, Union
from pydantic import BaseModel, Field
from fastapi import FastAPI, HTTPException, status
from sentence_transformers import SentenceTransformer

app = FastAPI(
    title="Fikra Embedding API",
    version="1.0.0",
    description="OpenAI-compatible local vector embedding service"
)

# Global state allocation for the verified model leader
MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
print(f"Loading production embedding vector space: {MODEL_NAME}...")
embedding_model = SentenceTransformer(MODEL_NAME)
print("Embedding model loaded and ready for inference pipeline.")

class EmbeddingRequest(BaseModel):
    input: Union[str, List[str]] = Field(..., description="The input text string or array of strings.")
    model: str = Field(default="fikra-embedding-v1", description="The model ID.")

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

class EmbeddingUsage(BaseModel):
    prompt_tokens: int = 0
    total_tokens: int = 0

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

@app.post("/v1/embeddings", response_model=EmbeddingResponse, status_code=status.HTTP_200_OK)
async def create_embeddings(request: EmbeddingRequest):
    try:
        if isinstance(request.input, str):
            input_data = [request.input]
        elif isinstance(request.input, list):
            input_data = request.input
            if not input_data:
                raise HTTPException(status_code=400, detail="The input array cannot be empty.")
        else:
            raise HTTPException(status_code=422, detail="Input must be a string or array of strings.")

        # Generate vectors
        embeddings = embedding_model.encode(input_data, normalize_embeddings=True, show_progress_bar=False)

        response_data = []
        estimated_tokens = 0
        
        for index, vector in enumerate(embeddings):
            response_data.append(EmbeddingData(index=index, embedding=vector.tolist()))
            estimated_tokens += max(1, len(input_data[index].split()))

        return EmbeddingResponse(
            data=response_data,
            model=request.model,
            usage=EmbeddingUsage(prompt_tokens=estimated_tokens, total_tokens=estimated_tokens)
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference Engine Exception: {str(e)}")

@app.get("/")
async def health_check():
    return {"status": "healthy", "model_loaded": MODEL_NAME, "timestamp": time.time()}