import os
import gradio as gr
import json
from gradio_client import Client, handle_file

backend = Client(os.getenv("BACKEND"), hf_token=os.getenv("TOKEN"))
JS_FUNC1 = os.getenv("JS_FUNC1")
JS_FUNC2 = os.getenv("JS_FUNC2")

def detect(image):
    try:
        file_1 = handle_file(image)
    except Exception as e:
        gr.Info("Please upload an image file.")
        return "", "", "", {}

    result_text = backend.predict(
        image=handle_file(image),
        api_name="/detect"
    )

    result = json.loads(result_text)
    if result and result["status"] == "ok":
        overall = result["overall"]
        aigen = result["aigen"]
        deepfake = result["deepfake"]

        breakdown = {
            "GenAI": 99,
            "Face Manipulation": 99,
            "Diffusion (Flux)": 99,
            "StyleGAN": 0,
        }
        return overall, aigen, deepfake, breakdown
    else:
        raise gr.Error("Error in processing image")

custom_css = """
.result-card {
  background: white;
  border-radius: 20px;
  padding: 24px;
  box-shadow: 0 6px 20px rgba(0,0,0,0.1);
}
.result-title {
  font-size: 20px;
  font-weight: bold;
  margin-bottom: 12px;
}
.progress-container {
  margin-bottom: 12px;
}
.progress-label {
  font-size: 14px;
  font-weight: 500;
}
.progress-bar {
  height: 12px;
  border-radius: 6px;
  background: #e5e7eb;
  margin-top: 4px;
  overflow: hidden;
}
.progress-fill {
  height: 100%;
  background: linear-gradient(90deg, #ff416c, #ff4b2b);
}
"""

MARKDOWN0 = """
# 🔍 DeepFake Detector
Upload an image and check if it’s AI-generated or manipulated.
"""
MARKDOWN3 = """
<div align="right"><a href="https://faceonlive.com/face-search-online" target='_blank'>Reverse Face Search</a></div>
<div align="right"><a href="https://faceonlive.com/reverse-image-search" target='_blank'>Reverse Image Search</a></div>
"""

def breakdown_ui(breakdown: dict):
    html = ""
    for label, val in breakdown.items():
        html += f"""
        <div class="progress-container">
            <div class="progress-label">{label}: {val}%</div>
            <div class="progress-bar"><div class="progress-fill" style="width:{val}%;"></div></div>
        </div>
        """
    return html

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown(MARKDOWN0)

    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(type='filepath', height=360, label="Upload Image")
            detect_button = gr.Button("🚀 Detect", elem_classes="button-gradient")
        with gr.Column(scale=2):
            overall = gr.Label(label="Overall Result")
            aigen = gr.Label(label="Generative AI Model")
            deepfake = gr.Label(label="Face Manipulation")
            breakdown_html = gr.HTML(label="Detection Breakdown")

    gr.Markdown(MARKDOWN3)

    detect_button.click(
        detect,
        inputs=image,
        outputs=[overall, aigen, deepfake, breakdown_html],
    ).then(
        breakdown_ui,
        inputs=[breakdown_html],
        outputs=breakdown_html
    )

demo.launch(server_name="0.0.0.0", share=True)