import gradio as gr
from PIL import Image
import os
import sys

current_dir = os.path.dirname(os.path.abspath(__file__))
defake_dir = os.path.join(current_dir, "defake")
if defake_dir not in sys.path:
    sys.path.insert(0, defake_dir)

import defake.test_api
# 把 defake.test_api 里的 NeuralNet 类，“挂载”到当前的主模块 (__main__) 上
setattr(sys.modules['__main__'], 'NeuralNet', defake.test_api.NeuralNet)

from app_infer import run_infer_from_image

APP_API_TOKEN = os.environ.get("APP_API_TOKEN", "")

def detect_api(image: Image.Image, api_token: str = ""):
    if APP_API_TOKEN and api_token != APP_API_TOKEN:
        raise gr.Error("Invalid API token.")
    result = run_infer_from_image(image)
    return result

def detect_ui(image: Image.Image):
    result = run_infer_from_image(image)
    return result["fake_score"], str(result["is_fake"])

with gr.Blocks(title="CISPA Citizen DeFake") as demo:
    gr.Markdown("# CISPA Citizen DeFake")
    gr.Markdown("Upload an image to detect if it is fake.")

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="Upload Image")
            btn = gr.Button("Detect")
        with gr.Column():
            score_out = gr.Number(label="Fake score (1=fake)")
            is_fake_out = gr.Textbox(label="Prediction (True=fake)")

    btn.click(fn=detect_ui, inputs=img_input, outputs=[score_out, is_fake_out])

    gr.Markdown("## API")
    gr.Markdown(
        "POST to `/run/predict` with form-data: `image`, `api_token`.\n"
        "Returns JSON: `{fake_score: float, is_fake: bool, probs: [p0,p1], pred_class: int}`."
    )

    api = gr.Interface(
        fn=detect_api,
        inputs=[
            gr.Image(type="pil", label="image"),
            gr.Textbox(label="api_token", type="password"),
        ],
        outputs="json",
        # allow_flagging="never",
        title="API Endpoint",
    )

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