# ml-service/main.py
from fastapi import FastAPI
from sentence_transformers import SentenceTransformer
from PIL import Image
from pydantic import BaseModel
from typing import Optional
import requests
from io import BytesIO
import uvicorn

app = FastAPI()
image_model = SentenceTransformer('clip-ViT-B-32')
text_model = SentenceTransformer('all-MiniLM-L6-v2')

class EmbedRequest(BaseModel):
    image_url: Optional[str] = None
    text_input: str

@app.get("/health")
def health():
    return {"status": "ok", "message": "ML service is running"}

@app.post("/embed")
def embed(req: EmbedRequest):
    text_embedding = text_model.encode(req.text_input).tolist()

    image_embedding = None
    if req.image_url:
        response = requests.get(req.image_url)
        img = Image.open(BytesIO(response.content)).convert('RGB')
        image_embedding = image_model.encode(img).tolist()

    return {
        "image_embedding": image_embedding,
        "text_embedding": text_embedding
    }

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="127.0.0.1",  # only accept connections from the same server, not the public internet
        port=8002,
        reload=False,       # reload should be False on a server, only useful during local dev
    )