from fastapi import FastAPI
from pymongo import MongoClient
from sentence_transformers import SentenceTransformer
from PIL import Image
from io import BytesIO
import numpy as np
import requests
from bson import ObjectId
import os

# ===============================================================
# 🚀 FastAPI App
# ===============================================================
app = FastAPI(title="Educational Placemat Embedding API")

# ===============================================================
# 🔐 Database Connection
# ===============================================================
MONGO_URI = "mongodb+srv://anna_db_user:6zxpOoyMUqnpxrBS@similaritysearch.xblvd4g.mongodb.net/"
client = MongoClient(MONGO_URI)
db = client["similaritysearch"]

# ===============================================================
# 🧩 Load Lightweight CLIP Model
# ===============================================================
os.makedirs("/tmp/model_cache", exist_ok=True)
model = SentenceTransformer(
    "sentence-transformers/clip-ViT-B-16",
    cache_folder="/tmp/model_cache"
)

# ===============================================================
# 🖼️ Endpoint — Generate Image Embedding
# ===============================================================
@app.post("/generate_embedding")
def generate_embedding(data: dict):
    """
    Generate a CLIP embedding for an image from URL and store it in MongoDB.
    """
    img_url = data["thumbnail"]
    image = Image.open(BytesIO(requests.get(img_url).content)).convert("RGB").resize((512, 512))
    emb = model.encode(image, convert_to_numpy=True, normalize_embeddings=True)
    db.images.update_one({"_id": ObjectId(data["_id"])}, {"$set": {"embedding": emb.tolist()}})
    return {"message": "✅ Embedding added successfully"}

# ===============================================================
# 🏠 Root route
# ===============================================================
@app.get("/")
def home():
    return {"status": "running", "model": "clip-ViT-B-16"}