import json
from datetime import datetime
from pathlib import Path
from uuid import uuid4

import gradio as gr
import numpy as np
from PIL import Image

from huggingface_hub import CommitScheduler, InferenceClient


IMAGE_DATASET_DIR = Path("image_dataset") / f"train-{uuid4()}"
IMAGE_DATASET_DIR.mkdir(parents=True, exist_ok=True)
IMAGE_JSONL_PATH = IMAGE_DATASET_DIR / "metadata.jsonl"

scheduler = CommitScheduler(
    repo_id="example-commit-scheduler-image",
    repo_type="dataset",
    folder_path=IMAGE_DATASET_DIR,
    path_in_repo=IMAGE_DATASET_DIR.name,
)

client = InferenceClient()


def generate_image(prompt: str) -> Image:
    return client.text_to_image(prompt)


def save_image(prompt: str, image_array: np.ndarray) -> None:
    image_path = IMAGE_DATASET_DIR / f"{uuid4()}.png"

    with scheduler.lock:
        Image.fromarray(image_array).save(image_path)
        with IMAGE_JSONL_PATH.open("a") as f:
            json.dump({"prompt": prompt, "file_name": image_path.name, "datetime": datetime.now().isoformat()}, f)
            f.write("\n")


def get_demo():
    with gr.Row():
        prompt_value = gr.Textbox(label="Prompt")
        image_value = gr.Image(label="Generated image")
    text_to_image_btn = gr.Button("Generate")
    text_to_image_btn.click(fn=generate_image, inputs=prompt_value, outputs=image_value).success(
        fn=save_image,
        inputs=[prompt_value, image_value],
        outputs=None,
    )