# main.py
from __future__ import annotations

import asyncio
import time
from typing import List, Optional, Dict, Any

from fastapi import FastAPI, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, validator

from get_embedding import EmbeddingFetcher


# -----------------------------
# Request / Response Schemas
# -----------------------------
class TextListRequest(BaseModel):
    texts: List[str] = Field(..., description="List of strings to embed", min_items=1)

    @validator("texts")
    def non_empty_texts(cls, v: List[str]) -> List[str]:
        if any((t is None or not isinstance(t, str) or t.strip() == "") for t in v):
            raise ValueError("All items in 'texts' must be non-empty strings.")
        return v


class EmbeddingResponse(BaseModel):
    model_id: str
    device: str
    dims: int
    count: int
    elapsed_ms: float
    embeddings: List[List[float]]


# -----------------------------
# App factory with lifespan
# -----------------------------
def create_app() -> FastAPI:
    app = FastAPI(title="Embedding API", version="1.0.0")

    # CORS: keep your original open policy (tighten in production)
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # Global container for services
    app.state.container: Dict[str, Any] = {}
    app.state.init_lock = asyncio.Lock()

    @app.on_event("startup")
    async def on_startup() -> None:
        # Initialize the EmbeddingFetcher once, asynchronously
        async with app.state.init_lock:
            if "embedder" not in app.state.container:
                fetcher = EmbeddingFetcher()
                # Build models / download snapshots off the main thread
                await fetcher.ensure_ready()
                app.state.container["embedder"] = fetcher

    @app.on_event("shutdown")
    async def on_shutdown() -> None:
        # Nothing special is required, but the hook is here for future cleanup
        pass

    # ---------------
    # Dependencies
    # ---------------
    def get_embedder() -> EmbeddingFetcher:
        fetcher: Optional[EmbeddingFetcher] = app.state.container.get("embedder")
        if fetcher is None:
            # Defensive: if a request sneaks in before startup finishes
            raise HTTPException(status_code=503, detail="Service not ready. Try again shortly.")
        return fetcher

    # ---------------
    # Routes
    # ---------------
    @app.get("/", tags=["meta"])
    async def home():
        return {"status": "ok", "message": "Embedding service is running."}

    @app.get("/healthz", tags=["meta"])
    async def healthz():
        # Lightweight health; could add a test encode if you want deeper checks
        return {"status": "healthy"}

    @app.post("/get-embedding/", response_model=EmbeddingResponse, tags=["embedding"])
    async def get_embedding(request: TextListRequest, embedder: EmbeddingFetcher = Depends(get_embedder)):
        # Offload embedding to the service (async wrapper over blocking HF/Torch calls)
        start = time.perf_counter()
        try:
            vectors = await embedder.embed(request.texts)
        except ValueError as ve:
            raise HTTPException(status_code=400, detail=str(ve)) from ve
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Embedding failed: {e}") from e

        elapsed_ms = (time.perf_counter() - start) * 1000.0
        dims = len(vectors[0]) if vectors and len(vectors[0]) else 0

        return EmbeddingResponse(
            model_id=embedder.model_id,
            device=embedder.device_str,
            dims=dims,
            count=len(vectors),
            elapsed_ms=round(elapsed_ms, 3),
            embeddings=vectors,
        )

    return app


app = create_app()

# Optional: run via `python main.py` in development
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,          # Turn off in production
        workers=1,            # Use a process manager (e.g., gunicorn) to scale
        log_level="info",
      )