import os
import uuid
import shutil
import numpy as np
import gradio as gr
from PIL import Image

from trufor_api import SplicingDetector

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
TRUFOR_FOLDER = os.path.join(SCRIPT_DIR, "TruFor_train_test")

detector = SplicingDetector(trufor_dir=TRUFOR_FOLDER, python_executable="python")


def predict(img: Image.Image):
    """Runs TruFor on an uploaded image and returns score + anomaly/confidence maps."""
    req_id = uuid.uuid4().hex  # unique per request

    uploads_dir = os.path.join(SCRIPT_DIR, "temp_uploads")
    out_root = os.path.join(SCRIPT_DIR, "temp_results")
    os.makedirs(uploads_dir, exist_ok=True)
    os.makedirs(out_root, exist_ok=True)

    image_path = os.path.join(uploads_dir, f"input_{req_id}.jpg")
    out_dir = os.path.join(out_root, f"run_{req_id}")
    os.makedirs(out_dir, exist_ok=True)

    img.save(image_path)

    try:
        results = detector.analyze_image(image_path, output_dir=out_dir)
        if results is None:
            raise gr.Error("TruFor failed. Check Space logs for the underlying error.")

        score = results["global_score"]

        anomaly = (np.clip(results["anomaly_map"], 0, 1) * 255).astype(np.uint8)
        conf = (np.clip(results["confidence_map"], 0, 1) * 255).astype(np.uint8)

        anomaly_img = Image.fromarray(anomaly)
        conf_img = Image.fromarray(conf)

        return score, anomaly_img, conf_img

    finally:
        # Cleanup (recommended for free Spaces)
        try:
            if os.path.exists(image_path):
                os.remove(image_path)
            shutil.rmtree(out_dir, ignore_errors=True)
        except Exception:
            pass


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Number(label="Global Score"),
        gr.Image(type="pil", label="Anomaly Map"),
        gr.Image(type="pil", label="Confidence Map"),
    ],
    title="TruFor Splicing Detector",
    description="Upload an image to get a manipulation score + heatmaps from TruFor.",
)

# IMPORTANT: launch must be at top-level (not inside if __name__ == '__main__')
demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False)