from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
import numpy as np
from sentence_transformers import SentenceTransformer
import asyncio
import logging
from typing import List, Optional, Union
from pydantic import BaseModel
import time

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Embedding API")

MODEL_NAME = "all-MiniLM-L12-v2"

# Load model on startup
try:
    model = SentenceTransformer(MODEL_NAME)
    logger.info("Model loaded successfully")
except Exception as e:
    logger.error(f"Failed to load model: {e}")
    model = None

# Request batching queue
class EmbeddingRequest(BaseModel):
    input: Union[str, List[str]]
    model: str = MODEL_NAME
    encoding_format: Optional[str] = None

class BatchItem:
    def __init__(self, texts: List[str]):
        self.texts = texts
        self.future: asyncio.Future = asyncio.Future()

pending_batch: List[BatchItem] = []
batch_lock = asyncio.Lock()
batch_event = asyncio.Event()

BATCH_TIMEOUT = 0.05  # 50ms window to collect requests
MAX_BATCH_SIZE = 32


async def batch_processor():
    """Continuously process batches of requests"""
    global pending_batch
    
    while True:
        try:
            # Wait for requests or timeout
            try:
                await asyncio.wait_for(batch_event.wait(), timeout=BATCH_TIMEOUT)
                batch_event.clear()
            except asyncio.TimeoutError:
                pass
            
            async with batch_lock:
                if not pending_batch:
                    continue
                
                # Take up to MAX_BATCH_SIZE items
                batch_to_process = pending_batch[:MAX_BATCH_SIZE]
                pending_batch = pending_batch[MAX_BATCH_SIZE:]
                
                if not batch_to_process:
                    continue
            
            # Flatten all texts
            all_texts = []
            text_counts = []
            for item in batch_to_process:
                all_texts.extend(item.texts)
                text_counts.append(len(item.texts))
            
            logger.info(f"Processing batch of {len(batch_to_process)} requests, {len(all_texts)} texts total")
            
            # Compute embeddings
            start = time.time()
            embeddings = model.encode(all_texts, convert_to_numpy=True)
            elapsed = time.time() - start
            logger.info(f"Embedding computed in {elapsed:.2f}s")
            
            # Split embeddings back to individual requests
            idx = 0
            for item, count in zip(batch_to_process, text_counts):
                item_embeddings = embeddings[idx:idx+count].tolist()
                item.future.set_result(item_embeddings)
                idx += count
        
        except Exception as e:
            logger.error(f"Error in batch processor: {e}")
            async with batch_lock:
                for item in pending_batch:
                    if not item.future.done():
                        item.future.set_exception(e)
                pending_batch.clear()
            await asyncio.sleep(1)


@app.on_event("startup")
async def startup():
    """Start batch processor on app startup"""
    if model is None:
        raise RuntimeError("Model failed to load")
    asyncio.create_task(batch_processor())
    logger.info("Batch processor started")


@app.get("/v1/models")
async def list_models():
    """OpenAI-compatible models endpoint"""
    if model is None:
        raise HTTPException(status_code=503, detail="Model not ready")
    
    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_NAME,
                "object": "model",
                "created": 0,
                "owned_by": "huggingface",
                "permission": [],
                "root": MODEL_NAME,
                "parent": None
            }
        ]
    }


@app.post("/v1/embeddings")
async def create_embeddings(request: EmbeddingRequest):
    if model is None:
        raise HTTPException(status_code=503, detail="Model not ready")
    
    # Normalize input to list
    if isinstance(request.input, str):
        texts = [request.input]
    else:
        texts = request.input
    
    if not texts:
        raise HTTPException(status_code=400, detail="input cannot be empty")
    
    if len(texts) > 512:
        raise HTTPException(status_code=400, detail="Cannot embed more than 512 texts at once")
    
    # Create batch item
    batch_item = BatchItem(texts)
    
    async with batch_lock:
        pending_batch.append(batch_item)
        
        # Signal processor if we hit batch size
        if len(pending_batch) >= MAX_BATCH_SIZE:
            batch_event.set()
    
    # Signal processor that there's work
    batch_event.set()
    
    # Wait for result with timeout
    try:
        embeddings_list = await asyncio.wait_for(batch_item.future, timeout=30.0)
    except asyncio.TimeoutError:
        raise HTTPException(status_code=504, detail="Embedding request timed out")
    
    # Format response as OpenAI-compatible
    data = []
    for idx, embedding in enumerate(embeddings_list):
        data.append({
            "object": "embedding",
            "embedding": embedding,
            "index": idx
        })
    
    return {
        "object": "list",
        "data": data,
        "model": MODEL_NAME,
        "usage": {
            "prompt_tokens": sum(len(text.split()) for text in texts),
            "total_tokens": sum(len(text.split()) for text in texts)
        }
    }


@app.get("/health")
async def health():
    """Health check endpoint"""
    if model is None:
        return JSONResponse({"status": "unhealthy", "reason": "model not loaded"}, status_code=503)
    return {"status": "healthy"}


@app.get("/")
async def root():
    """API info"""
    return {
        "name": "OpenAI-Compatible Embedding API",
        "model": MODEL_NAME,
        "endpoints": {
            "GET /v1/models": "List available models",
            "POST /v1/embeddings": "Create embeddings (OpenAI-compatible)",
            "GET /health": "Health check"
        }
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)