from PIL import Image
import gradio as gr
import io, os, tempfile

def _get_path(uploaded):
    if uploaded is None:
        return None
    if isinstance(uploaded, dict):
        return uploaded.get("name")
    try:
        return uploaded.name
    except Exception:
        return uploaded

def process(uploaded, keep_metadata):
    """uploaded: gr.File input; keep_metadata: checkbox (opt-in)
    Returns: preview (PIL.Image), downloadable file path, status message
    """
    path = _get_path(uploaded)
    if not path:
        return None, None, "No file uploaded."
    try:
        img = Image.open(path)
    except Exception as e:
        return None, None, f"Cannot open image: {e}"

    fmt = (img.format or "PNG").upper()

    if keep_metadata:
        # Opt-in path: return original bytes so metadata is preserved
        with open(path, "rb") as f:
            data = f.read()
        preview = img.convert("RGB")
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.' + fmt.lower())
        tmp.write(data); tmp.flush(); tmp.close()
        return preview, tmp.name, "Original metadata preserved (opt‑in)."

    # Privacy path: re-save into a fresh buffer (no exif passed)
    preview_img = img.convert("RGB")
    buf = io.BytesIO()
    preview_img.save(buf, format=fmt)
    buf.seek(0)

    # Verify removal: try getexif(); fallback to info['exif']
    try:
        img2 = Image.open(buf)
        ex = img2.getexif()
        has_exif = bool(len(ex))
    except Exception:
        has_exif = bool(img2.info.get("exif")) if hasattr(img2, 'info') else False

    ok = not has_exif
    msg = "Metadata removed ✓" if ok else "Metadata may remain (verification failed)."

    # write sanitized file for download
    out_tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.' + fmt.lower())
    out_tmp.write(buf.getvalue()); out_tmp.flush(); out_tmp.close()

    return preview_img, out_tmp.name, msg

with gr.Blocks() as demo:
    gr.Markdown("**Upload an image — metadata will be stripped unless you opt into keeping it.**")
    file_in = gr.File(label='Upload image (PNG/JPEG...)')
    keep = gr.Checkbox(label='Keep original metadata (opt‑in)', value=False)
    img_out = gr.Image(label='Processed preview')
    file_out = gr.File(label='Download processed file')
    status = gr.Textbox(label='Status')
    btn = gr.Button('Process')
    btn.click(process, inputs=[file_in, keep], outputs=[img_out, file_out, status])
    demo.launch()