import torch
from torch import nn
from transformers import AutoImageProcessor, Dinov2Model
from PIL import Image
import base64
from io import BytesIO
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import imagehash  # <-- NOVO

MODEL_NAME = "facebook/dinov2-small"
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
model = Dinov2Model.from_pretrained(MODEL_NAME)
projection = nn.Linear(model.config.hidden_size, 512)

app = FastAPI(
    title="API de Embedding de Imagem",
    description="Endpoint para obter o embedding e pHash de uma imagem.",
    version="1.0.0"
)

class ImageRequest(BaseModel):
    image: str

@app.post("/embed")
async def get_embedding(request: ImageRequest):
    try:
        header, img_base64 = request.image.split(",", 1)
        image_data = base64.b64decode(img_base64)
        image = Image.open(BytesIO(image_data)).convert("RGB")
        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
            last_hidden_state = outputs.last_hidden_state
            embedding = last_hidden_state[:, 0]
            embedding_512 = projection(embedding)
        phash = str(imagehash.phash(image))  # <-- NOVO
        return {
            "embedding": embedding_512.squeeze().tolist(),
            "phash": phash
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Erro ao processar a imagem: {e}")