import json
import chromadb
from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
from chromadb.utils.data_loaders import ImageLoader
import ollama
import base64
import io
from io import BytesIO
import openai
from google import genai
import os 
from PIL import Image
import numpy as np
from google.genai.types import Part, Content

client = genai.Client(api_key=os.environ.get("GENAI_API_KEY"))

class PlatinumPipeline:

    def __init__(self, json_file="pricelist.json"):
        self.json_file = json_file
        self.collection_name = "pricetag"
        self.TOP_K = 3
        self.chroma_client = chromadb.Client()
        self.image_loader = ImageLoader()
        self.multimodal_ef = OpenCLIPEmbeddingFunction()
        self.model_name = "gemma3:4b"
        self.multimodal_db = self.chroma_client.create_collection(
            name=self.collection_name,
            embedding_function=self.multimodal_ef,
            data_loader=self.image_loader
        )
        with open(self.json_file, "r") as f:
            self.items = json.load(f)

        self.valid_items, self.valid_ids, self.valid_uris, self.valid_metadatas = [], [], [], []

        for i, item in enumerate(self.items):
            image_path = item["image"]
            try:
                with open(image_path, "rb"): 
                    self.valid_items.append(item)
                    self.valid_ids.append(str(i))
                    self.valid_uris.append(image_path)
                    self.valid_metadatas.append({
                        "name": item["item_name"],
                        "price": item["price"],
                        "original_index": i
                    })
            except FileNotFoundError:
                continue

        if self.valid_items:
            try:
                self.multimodal_db.add(
                    ids=self.valid_ids,
                    uris=self.valid_uris,
                    metadatas=self.valid_metadatas
                )
            except Exception as e:
                print(e)
        else:
            print("Records are empty")

    def query(self, image_bytes: bytes):
        context = []
        if self.multimodal_db.count() == 0:
            return "No records found in database!!!"

        try:
            pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
            image_np = np.array(pil_image) 
            
            query_results = self.multimodal_db.query(
                query_images=[image_np],  
                n_results=min(self.TOP_K, self.multimodal_db.count()),
                include=['uris', 'metadatas', 'distances', 'data']
            )

            if query_results['ids'] and len(query_results['ids'][0]) > 0:
                for i in range(len(query_results['ids'][0])):
                    context.append(query_results['metadatas'][0][i])
                
               
                prompt = (
                    f"Context:\n{context}\n\n"
                    "You have to write the estimated price of the item in the image "
                    "that has been provided based on the context which contains "
                    "the names and prices of similar items in the database."
                )

                text_part = Part(text=prompt)
                image_part = Part(
                    inline_data={
                        "mime_type": "image/jpeg",
                        "data": base64.b64encode(image_bytes).decode('utf-8')
                    }
                )
            
                content = Content(parts=[text_part, image_part])
                response = client.models.generate_content(
                    model="gemini-2.0-flash-exp",  
                    contents=content
                )
                
                return response.text
            else:
                return "No context found"
        except Exception as e:
            return f"Error: {str(e)}"