import io
import uvicorn
from fastapi import FastAPI, UploadFile, status, Form
from contextlib import asynccontextmanager
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import BaseModel
from config import settings
from PIL import Image
from faiss_db import FAISS_DB
from vlm_model import VisionLanguageModel
import torch
import gc
import base64


vectordb = vectordb = FAISS_DB(
    faiss_index_file_path=settings.FAISS_INDEX_FILE_PATH,
    keyframes_groups_json_path=settings.KEYFRAMES_GROUPS_JSON_PATH,
)
vlm_model = VisionLanguageModel(
    model_name=settings.MODEL_NAME,
    pretrained=settings.MODEL_CHECKPOINT,
    device=(
        "cuda" if settings.DEVICE == "cuda" and torch.cuda.is_available() else "cpu"
    ),
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    global vectordb
    global vlm_model
    yield
    del vectordb
    del vlm_model
    torch.cuda.empty_cache()
    gc.collect()
    torch.cuda.ipc_collect()


app = FastAPI(title="CLIP Retrieval API", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_headers=["*"],
    allow_credentials=True,
    allow_methods=["*"],
)


@app.get("/", include_in_schema=False)
async def root() -> None:
    return RedirectResponse("/docs")


@app.get("/health", status_code=status.HTTP_200_OK, tags=["health"])
async def perform_healthcheck() -> None:
    return JSONResponse(content={"message": "success"})


class Item(BaseModel):
    query: str | UploadFile
    queryType: str
    topk: int


@app.post("/retrieval")
async def retrieval(
    query: str | UploadFile = Form(...),
    queryType: str = Form("text"),  # text || image
    topk: int = Form(...),
) -> JSONResponse:
    try:
        if queryType == "text":
            input_query = query
        elif queryType == "image":
            image_bytes = await query.read()
            input_query = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        else:
            return JSONResponse(
                status_code=status.HTTP_400_BAD_REQUEST,
                content={
                    "message": "Invalid query type. Provide either a text query or an image."
                },
            )

        query_vector = vlm_model.get_embedding(input_query)
    except Exception as e:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={"message": f"Invalid query: {str(e)}"},
        )

    try:
        search_results = vectordb.kNN_search(query_vector, topk)
    except Exception as e:
        return JSONResponse(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            content={"message": f"Search error: {str(e)}"},
        )

    return JSONResponse(
        status_code=status.HTTP_200_OK,
        content={"message": "success", "details": search_results},
    )


# Start API
if __name__ == "__main__":
    uvicorn.run("main:app", host=settings.HOST, port=settings.PORT, reload=True)