import os
import vertexai
from vertexai.vision_models import Image, MultiModalEmbeddingModel
import base64
import io
from PIL import Image as PILImage
from typing import List

class VertexVisionService:
    def __init__(self, project_id: str, location: str):
        # Ensure credentials are set before init if using service account key
        # (This is now handled globally in main.py)
        vertexai.init(project=project_id, location=location)
        self.model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding@001")

    def get_image_embedding(self, image_path: str) -> List[float]:
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image not found at {image_path}")
        image = Image.load_from_file(image_path)
        # Standardizing to 1408 dimensions for consistency
        embeddings = self.model.get_embeddings(image=image, dimension=1408)
        return embeddings.image_embedding

    def get_text_embedding(self, text: str) -> List[float]:
        """Converts text query into a 1408-d vector."""
        embeddings = self.model.get_embeddings(
            contextual_text=text, 
            dimension=1408
        )
        return embeddings.text_embedding

    def encode_image_to_base64(self, image_path: str) -> str:
        """Resizes image to 640x640 and encodes to Base64."""
        with PILImage.open(image_path) as img:
            img.thumbnail((640, 640))
            buffered = io.BytesIO()
            img.save(buffered, format="JPEG", quality=85)
            return base64.b64encode(buffered.getvalue()).decode('utf-8')
