import gradio as gr
from transformers import pipeline

from glob import glob
from PIL import Image
import os
from icrawler.builtin import GoogleImageCrawler

def download_image(query, out_file):
  google_crawler = GoogleImageCrawler(storage={'root_dir': './'})
  google_crawler.crawl(keyword=query, max_num=1, overwrite=True)
  os.rename(glob("./000001.*")[0], out_file)

def generate_story(prompt):
  story = storygen(f"{prompt}")[0]['generated_text']
  return story

def start_neural_style_transfer(img1_name, img2_name):
  img1_filename, img2_filename = "style_transfer_1.jpg", "style_transfer_2.jpg"

  download_image(img1_name, img1_filename)
  download_image(img2_name, img2_filename)

  styled_image = nst(img1_filename, img2_filename)

  pil_img = Image.open(styled_image)

  return pil_img

def detect_objects(file_name):
  out_img = detectron(file_name)
  pil_img = Image.open(out_img) 
  return pil_img

def main(text_input):
  text_output, image_output, metadata = None, None, None 

  task_type_q = f"User: {text_input}\nWhat is the task the user is asking to do?\n \
                - story generation task\n \
                - image style transfer task\n \
                - object detection task"
  task = t0pp(task_type_q)
  task = task.lower().replace('.', '')

  if task=="story generation task":
    story_prompt = t0pp(f"User: {text_input}\nWhat story is the user asking for?")
    text_output = generate_story(story_prompt)
    metadata = f"Prompt used to generate the story:\n{story_prompt}"

  elif task=="image style transfer task":
    img1_name = t0pp(f"User: {text_input}\nWhat is the name of the picture to which style is to be tranferred?")
    img2_name = t0pp(f"User: {text_input}\nWhat is the name of the picture from which style is to be tranferred?")
    image_output = start_neural_style_transfer(img1_name, img2_name)
    metadata = f"Image from which style is to be transferred: {img2_name}\nImage to which style is to be transferred: {img1_name}"

  elif task=="object detection task":
    img_file = "object_detection.jpg"
    img_name = t0pp(f"User: {text_input}\nWhat image is the user referring to?")
    download_image(img_name, img_file)
    image_output = detect_objects(img_file)
    metadata = f"Image from which objects are to be detected: {img_name}"

  return text_output, image_output

if __name__=="__main__":
    t0pp = gr.Interface.load("huggingface/bigscience/T0pp")
    storygen = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
    nst = gr.Interface.load("spaces/luca-martial/neural-style-transfer")
    detectron = gr.Interface.load("spaces/akhaliq/Detectron2")

    title = "Multipurpose AI"
    description = "With this demo you can do image style transfer, story generation and object detection(as of now, detects all the possible objects in the image) using a single text box. Use the examples provided below to get started."
    article = "[Bigscience T0pp](https://huggingface.co/bigscience/T0pp) | [story generator](https://huggingface.co/pranavpsv/genre-story-generator-v2) | [style transfer](https://huggingface.co/spaces/luca-martial/neural-style-transfer) | [object detection](https://huggingface.co/spaces/akhaliq/Detectron2)"
    examples = [["Find an image of taj mahal and transfer the style of starry night to it"],
                ["Find an image of a football match and detect the objects in that image"],
                ["Give me a story about a hard working farmer"]]

    gr.Interface(
      main,
      title=title,
      description=description,
      article=article,
      inputs=gr.inputs.Textbox(lines=5, label="Input"),
      outputs=[gr.outputs.Textbox(label="Output"), gr.outputs.Image(label="Ouptut"),],
      examples=examples,
      enable_queue=True,
  ).launch(debug=True)