import requests
import zipfile
import gradio as gr

from src.examples import audio_examples_tab, examples_tab
from src.mk_attacks_variations import mk_audio_variations, mk_image_variations
from src.mk_leaderboard import mk_leaderboard

from pathlib import Path


abs_path = Path(__file__).parent

with gr.Blocks(theme=gr.themes.Base()) as demo:

    gr.Markdown(
        """
        # 🥇 Omni Seal Bench Watermarking Leaderboard
        """
    )
    with gr.Tabs():
        with gr.Tab("Audio"):
            gr.Markdown(
                """
                ### Performance on Ravdess dataset
                """
            )
            mk_leaderboard(
                abs_path / "data/audio_benchmark.csv",
                default_selection=[
                    "TimeDomain_bit_acc",
                    "AmplitudeDomain_bit_acc",
                    "identity_snr",
                    "identity_bit_acc",
                    "identity_detect_prob",
                    "avg_bit_acc",
                    "avg_tn_bit_acc",
                    "avg_detect_prob",
                    "avg_tn_detect_prob",
                ],
                core_columns=["model", "snr"],
                filter_columns=["model"],
                search_columns=["model"],
                categories={
                    "speed": "Time",
                    "updownresample": "Time",
                    "echo": "Time",
                    "random_noise": "Amplitude",
                    "lowpass_filter": "Amplitude",
                    "highpass_filter": "Amplitude",
                    "bandpass_filter": "Amplitude",
                    "smooth": "Amplitude",
                    "boost_audio": "Amplitude",
                    "duck_audio": "Amplitude",
                    "shush": "Amplitude",
                },
            )

            mk_audio_variations(abs_path / "data/audio_attacks_variations.csv")

        with gr.Tab("Image"):
            gr.Markdown(
                """
                ### Performance on Val2014 dataset
                """
            )
            mk_leaderboard(
                abs_path / "data/image_benchmark.csv",
                default_selection=[
                    "Visual_bit_acc",
                    "Geometric_bit_acc",
                    "Compression_bit_acc",
                    "Inpainting_bit_acc",
                    "Mixed_bit_acc",
                    "avg_bit_acc",
                    "avg_p_value",
                    "avg_word_acc",
                ],
                core_columns=[
                    "model",
                    "psnr",
                    "ssim",
                    "lpips",
                ],
                filter_columns=[
                    "model",
                ],
                search_columns=["model"],
                categories={
                    "proportion": "Geometric",
                    "collage": "Inpainting",
                    "crop": "Geometric",
                    "rot": "Geometric",
                    "jpeg": "Compression",
                    "brightness": "Visual",
                    "contrast": "Visual",
                    "saturation": "Visual",
                    "sharpness": "Visual",
                    "resize": "Geometric",
                    "overlay_text": "Inpainting",
                    "hflip": "Geometric",
                    "perspective": "Geometric",
                    "median_filter": "Visual",
                    "hue": "Visual",
                    "gaussian_blur": "Visual",
                    "comb": "Mixed",
                    "avg": "Averages",
                    "none": "Baseline",
                },
            )

            mk_image_variations(abs_path / "data/image_attacks_variations.csv")

        with gr.Tab("Image examples"):
            examples_tab(abs_path)

        with gr.Tab("Audio examples"):
            audio_examples_tab(abs_path)

        with gr.Tab("Docs"):
            README_URL = "https://raw.githubusercontent.com/facebookresearch/omnisealbench/refs/heads/main/README.md"

            def fetch_readme():
                response = requests.get(README_URL, timeout=4)
                if response.status_code == 200:
                    return response.text
                else:
                    return "Failed to fetch README.md. Please check the URL or try again later."

            # Define the Gradio interface
            gr.Markdown(fetch_readme())


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