import requests
import gradio as gr
from PIL import Image
import time

URL = "https://source.unsplash.com/random/500x500/?nature,fruit"

js_pos_select = """
async (images_selected_state) => {
  const gallery = document.querySelector("#gallery_id")
  const buttons_thumbnails = gallery.querySelectorAll(".thumbnails > button");
  const buttons_large = gallery.querySelectorAll(".grid-container > button");
  buttons_thumbnails.forEach((btn, idx) => {
    if(images_selected_state.includes(idx)){
      btn.classList.add('selected-custom');
    }else{
      btn.classList.remove('selected-custom');
    }
  })
  buttons_large.forEach((btn, idx) => {
    if(images_selected_state.includes(idx)){
      btn.classList.add('selected-custom');
    }else{
      btn.classList.remove('selected-custom');
    }
  })
  return images_selected_state
}
"""
css = """
.selected-custom {
    --ring-color: red !important;
    transform: scale(0.9) !important;
    border-color: red !important;
}
"""
def get_random_images(n=8):
  images = []
  for i in range(n):
    images.append(Image.open(requests.get(URL + f"?{time.time()}", stream=True).raw))
  return images

def get_select_index(imgs_set, evt: gr.SelectData):
  print(evt.index)
  if(evt.index in imgs_set):
    imgs_set.pop(imgs_set.index(evt.index))
  else:
    imgs_set.append(evt.index)
  print(imgs_set)
  return imgs_set

with gr.Blocks(css=css) as demo:
    images_selected_state = gr.JSON([], visible=False)

    with gr.Column(variant="panel"):
        with gr.Row(variant="compact"):
            btn = gr.Button("Generate image").style(full_width=False)

        gallery_photos = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery_id"
        ).style(columns=[4], rows=[2], object_fit="contain", height="auto")

    gallery_photos.select(
        get_select_index, images_selected_state, images_selected_state).then(
        None, images_selected_state, None, _js=js_pos_select)
    demo.load(get_random_images, None, gallery_photos)

demo.launch(debug=True)