import gradio as gr
from Clustering import ClusteringData
import numpy as np
from PIL import Image
import requests
import tempfile
import os
import logging
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


cd = ClusteringData()
cd.load_model_data()
logger.info("Clustering data loaded")


def search_images(text_query, uploaded_image, search_mode, top_k):
    preview = None
    results = []

    if search_mode == "Text" and text_query.strip():
        response = requests.get(
            f"https://ashish-001-clip-api.hf.space/embedding?text={text_query.strip()}")
        if response.status_code == 200:
            logger.info("Embedding returned successfully by text API")
            data = json.loads(response.content)
            embedding = data["embedding"]
            results = cd.find_similar_records(embedding, k=top_k)
        else:
            logger.info(f"{response.status_code} returned by the text API")
            results = []
        results = [os.path.join("coco", "val2017", "val2017", fname)
                   for i, fname in enumerate(results)]
        return None, results

    elif search_mode == "Image":
        if uploaded_image is not None:
            preview = uploaded_image
            tmp_path = uploaded_image
            # with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
            #     uploaded_image.save(tmp_file.name)
            #     tmp_path = tmp_file.name
        else:
            preview = 'Image.jpg'
            tmp_path = 'Image.jpg'
        url = "https://ashish-001-clip-api.hf.space/clip/process"
        files = {"file": open(tmp_path, "rb")}
        response = requests.post(url, files=files)
        if response.status_code == 200:
            embedding = np.array(response.json()['embedding']).squeeze()
            logger.info("Embedding returned successfully by image API")
            results = cd.find_similar_records(embedding, k=top_k)
        else:
            logger.info(
                f"{response.status_code} returned by the image API")
            results = []
        results = [os.path.join("coco", "val2017", "val2017", fname)
                   for i, fname in enumerate(results)]

        return preview, results


with gr.Blocks() as demo:
    gr.Markdown("## Multimodal Image Search with CLIP")
    gr.Markdown("Search images using **text** or **image upload**.")

    with gr.Row():
        with gr.Column(scale=1):
            # Inputs
            search_mode = gr.Radio(
                ["Text", "Image"], label="Search Mode", value="Text")
            text_input = gr.Textbox(
                label="Enter text query", placeholder="Type something...", visible=True, value='Empty street')
            file_input = gr.Image(
                type="filepath",
                label="Upload image",
                value="Image.jpg",
                visible=False
            )
            top_k = gr.Slider(1, 20, value=6, step=1,
                              label="Number of results")
            submit_btn = gr.Button("Search")

        with gr.Column(scale=2):
            preview_img = gr.Image(label="Uploaded / Default Image")
            result_gallery = gr.Gallery(
                label="Results", columns=3, height="auto")

    def toggle_inputs(mode):
        if mode == "Text":
            return (
                gr.update(visible=True),
                gr.update(visible=False, value=None),
                [],
                None
            )
        else:
            return (
                gr.update(visible=False),
                gr.update(visible=True, value=None),
                [],
                "Image.jpg"
            )

    search_mode.change(toggle_inputs, inputs=search_mode,
                       outputs=[text_input, file_input, result_gallery, preview_img])

    submit_btn.click(fn=search_images,
                     inputs=[text_input,
                             file_input, search_mode, top_k],
                     outputs=[preview_img, result_gallery,])

demo.launch()