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

import gradio as gr
from PIL import Image

from huggingface_hub import CommitScheduler

# ---------------- Dataset setup ----------------
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"

# Read Hugging Face token from secret
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN not found! Please set it in HF Spaces Secrets.")

# Scheduler for your HF dataset repo
scheduler = CommitScheduler(
    repo_id="codewithRiz/Buck_data_saving",  # your dataset repo
    repo_type="dataset",
    folder_path=IMAGE_DATASET_DIR,
    path_in_repo=IMAGE_DATASET_DIR.name,
    token=HF_TOKEN,  # use the token from secret
)

# ---------------- Save uploaded image ----------------
def save_uploaded_image(user_id: str, image: Image.Image) -> str:
    """
    Save the uploaded image to local dataset folder, log metadata, and push to HF Hub.
    """
    image_path = IMAGE_DATASET_DIR / f"{uuid4()}.png"

    with scheduler.lock:
        # Save the image locally
        image.save(image_path)

        # Save metadata to JSONL
        with IMAGE_JSONL_PATH.open("a") as f:
            json.dump({
                "user_id": user_id,
                "file_name": image_path.name,
                "datetime": datetime.now().isoformat()
            }, f)
            f.write("\n")

        # Automatically commit & push to HF Hub
        scheduler.commit(message=f"Add image {image_path.name} for user {user_id}")

    return f"Image uploaded and pushed to repo successfully for user {user_id}!"

# ---------------- Gradio UI ----------------
def get_demo():
    with gr.Row():
        user_id_input = gr.Textbox(label="User ID")
        image_input = gr.Image(label="Upload Image", type="pil")
    upload_btn = gr.Button("Upload")
    upload_status = gr.Textbox(label="Status")
    upload_btn.click(
        fn=save_uploaded_image,
        inputs=[user_id_input, image_input],
        outputs=upload_status
    )

# Launch Gradio app
demo = gr.Blocks()
with demo:
    get_demo()

demo.launch()