"""CLIP embedding service for frigate-notify.

Loads OpenAI ViT-L/14 CLIP model and exposes endpoints for generating
image embeddings and computing cosine similarity between embeddings.
Designed to run on a GPU machine (RTX 3090) as a sidecar service.
"""

import base64
import io
import logging
import time
from contextlib import asynccontextmanager

import numpy as np
import open_clip
import torch
from fastapi import FastAPI, HTTPException
from PIL import Image
from pydantic import BaseModel

logger = logging.getLogger("clip-service")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

# --- Pydantic models ---


class EmbedRequest(BaseModel):
    image: str  # base64-encoded image


class EmbedResponse(BaseModel):
    embedding: list[float]
    inference_ms: float


class SimilarityRequest(BaseModel):
    embedding1: list[float]
    embedding2: list[float]


class SimilarityResponse(BaseModel):
    score: float


class BatchSimilarityRequest(BaseModel):
    query_embedding: list[float]
    candidate_embeddings: list[list[float]]


class BatchSimilarityResponse(BaseModel):
    scores: list[float]
    best_index: int
    best_score: float


# --- Globals ---

model = None
preprocess = None
device = None
model_name = "ViT-L-14"
pretrained = "openai"


def load_model():
    """Load the CLIP model onto the best available device."""
    global model, preprocess, device

    if torch.cuda.is_available():
        device = torch.device("cuda")
        logger.info("CUDA available — using GPU: %s", torch.cuda.get_device_name(0))
    else:
        device = torch.device("cpu")
        logger.warning("CUDA not available — falling back to CPU")

    logger.info("Loading CLIP model %s (pretrained=%s)...", model_name, pretrained)
    t0 = time.monotonic()
    model, _, preprocess = open_clip.create_model_and_transforms(
        model_name, pretrained=pretrained, device=device
    )
    model.eval()
    elapsed = time.monotonic() - t0
    logger.info("Model loaded in %.1fs on %s", elapsed, device)


@asynccontextmanager
async def lifespan(app: FastAPI):
    load_model()
    yield


app = FastAPI(title="CLIP Embedding Service", version="1.0.0", lifespan=lifespan)


# --- Utility functions ---


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Compute cosine similarity between two vectors."""
    dot = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(dot / (norm_a * norm_b))


# --- Endpoints ---


@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "device": str(device) if device else "not_loaded",
        "model": f"{model_name}/{pretrained}",
    }


@app.post("/embed", response_model=EmbedResponse)
async def embed(request: EmbedRequest):
    if model is None or preprocess is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    # Decode base64 image
    try:
        image_bytes = base64.b64decode(request.image)
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid image data: {e}")

    # Preprocess and encode
    try:
        image_tensor = preprocess(image).unsqueeze(0).to(device)

        t0 = time.monotonic()
        with torch.no_grad():
            embedding = model.encode_image(image_tensor)
        inference_ms = (time.monotonic() - t0) * 1000

        # Normalize the embedding
        embedding = embedding / embedding.norm(dim=-1, keepdim=True)
        embedding_list = embedding.squeeze(0).cpu().numpy().tolist()

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference failed: {e}")

    return EmbedResponse(embedding=embedding_list, inference_ms=round(inference_ms, 2))


@app.post("/similarity", response_model=SimilarityResponse)
async def similarity(request: SimilarityRequest):
    a = np.array(request.embedding1, dtype=np.float32)
    b = np.array(request.embedding2, dtype=np.float32)

    if a.shape != b.shape:
        raise HTTPException(
            status_code=400,
            detail=f"Embedding dimensions must match: {a.shape} vs {b.shape}",
        )

    score = cosine_similarity(a, b)
    return SimilarityResponse(score=round(score, 6))


@app.post("/batch-similarity", response_model=BatchSimilarityResponse)
async def batch_similarity(request: BatchSimilarityRequest):
    query = np.array(request.query_embedding, dtype=np.float32)

    if not request.candidate_embeddings:
        raise HTTPException(status_code=400, detail="No candidate embeddings provided")

    candidates = np.array(request.candidate_embeddings, dtype=np.float32)

    if query.shape[0] != candidates.shape[1]:
        raise HTTPException(
            status_code=400,
            detail=f"Dimension mismatch: query={query.shape[0]}, candidates={candidates.shape[1]}",
        )

    # Compute cosine similarity for all candidates at once
    query_norm = query / (np.linalg.norm(query) + 1e-10)
    candidate_norms = candidates / (
        np.linalg.norm(candidates, axis=1, keepdims=True) + 1e-10
    )
    scores = candidate_norms @ query_norm

    scores_list = [round(float(s), 6) for s in scores]
    best_index = int(np.argmax(scores))
    best_score = scores_list[best_index]

    return BatchSimilarityResponse(
        scores=scores_list, best_index=best_index, best_score=best_score
    )