import os
import time
import uuid
from typing import Optional

import gradio as gr
from PIL import Image

# ------------------------------------------------------------
# Pick a writable save directory automatically.
# Order:
# 1) SAVE_DIR env var (if supplied)
# 2) /data/uploads           (when Persistent Storage is enabled)
# 3) /home/user/app/uploads  (always writable in Spaces)
# 4) /tmp/uploads            (always writable, ephemeral)
# ------------------------------------------------------------
def pick_writable_dir() -> str:
    candidates = []
    env_dir = os.getenv("SAVE_DIR")
    if env_dir:
        candidates.append(env_dir)
    candidates += ["/data/uploads", "/home/user/app/uploads", "/tmp/uploads"]

    for path in candidates:
        try:
            os.makedirs(path, exist_ok=True)
            # double-check we can create files here
            testfile = os.path.join(path, f".write_test_{uuid.uuid4().hex}")
            with open(testfile, "w") as f:
                f.write("ok")
            os.remove(testfile)
            return path
        except Exception:
            continue
    # last resort: current directory
    fallback = "./uploads"
    os.makedirs(fallback, exist_ok=True)
    return fallback

SAVE_DIR = pick_writable_dir()

ALLOWED_FORMATS = {"JPEG", "PNG"}  # strict type enforcement

def handle_upload(file_path: Optional[str]):
    if not file_path:
        return "Please choose an image.", None

    try:
        with Image.open(file_path) as im:
            fmt = (im.format or "").upper()
            if fmt not in ALLOWED_FORMATS:
                return f"Only JPEG/PNG allowed. Detected: {fmt or 'unknown'}.", None

            # Normalize to RGB so we can always save safely
            if im.mode not in ("RGB", "RGBA"):
                im = im.convert("RGB")

            # Choose extension from the detected format
            ext = ".jpg" if fmt == "JPEG" else ".png"
            filename = f"{int(time.time())}_{uuid.uuid4().hex}{ext}"
            save_path = os.path.join(SAVE_DIR, filename)

            # Reasonable defaults for each format
            if fmt == "JPEG":
                im.save(save_path, quality=90, optimize=True)
            else:
                im.save(save_path, optimize=True)

            info = (
                f"✅ Saved to **{save_path}**\n\n"
                f"- File: **{filename}**\n"
                f"- Size: **{im.width}×{im.height}px**\n"
                f"- Format: **{fmt}**"
            )
            return info, im.copy()
    except Exception as e:
        return f"Failed to read/save image: {e}", None

with gr.Blocks(title="Image Upload (JPEG/PNG)") as demo:
    gr.Markdown("### Upload a JPEG or PNG. The file will be validated and saved.")
    with gr.Row():
        inp = gr.Image(
            label="Drop or click to upload",
            type="filepath",              # keep original temp filepath
            image_mode="RGB",
            sources=["upload", "clipboard"]
        )
        preview = gr.Image(label="Preview", interactive=False)

    msg = gr.Markdown()

    # Auto-save when the image changes; also allow manual click.
    inp.change(handle_upload, inputs=inp, outputs=[msg, preview])
    gr.Button("Save").click(handle_upload, inputs=inp, outputs=[msg, preview])

    gr.Markdown(f"**Save directory:** `{SAVE_DIR}`")

if __name__ == "__main__":
    demo.launch()