import gradio as gr

from .expert_malware_detector import  expert_detect_malware
from .discriminative_chatter_detector import DiscriminativeChatterDetector
from .generative_malware_detector import GenerativeMalwareDetector


detector_a = DiscriminativeChatterDetector(dataset_name='blackbriar')
detector_c = GenerativeMalwareDetector()


def detect_malware_a(code: str, dependencies: str) -> str:
    """Model A detection"""
    result = detector_a.predict(code)
    return f"**Prediction:** {result['label']}"


def detect_malware_b(code: str, dependencies: str) -> str:
    """Model B detection"""
    result = expert_detect_malware(source_code=code, dependencies=dependencies)
    return f"**Prediction:** {result}"


def detect_malware_c(code: str, dependencies: str) -> str:
    """Model C detection"""
    result = detector_c.detect(code, dependencies)
    return f"**Prediction:** {result['label']}"


def clear_all_outputs():
    """
    Source - https://stackoverflow.com/a
    Posted by Vaishnav
    Retrieved 2025-11-17, License - CC BY-SA 4.0
    """
    return None, None, None, None, None


# hwr_theme = gr.themes.Default(primary_hue=gr.themes.colors.rose, secondary_hue=gr.themes.colors.red)
# with gr.Blocks(theme=hwr_theme) as demo:

with gr.Blocks() as demo:

    gr.Markdown("""
    # Malware Detection Demo

    Test three different AI models for detecting malicious code.
    One demo uses a rules-based model, one uses predictive ML ('discriminative') and one uses generative ML.

    Can you guess which is which?

    **Your task:** Try various inputs and track the responses to guess which demo is using which model type.
    """)

    with gr.Row():
        code_text = gr.Textbox(
            label="Enter source code to analyze",
            placeholder="Add your code here...",
            lines=3
        )
        dependencies_text = gr.Textbox(
            label="Enter dependency specificiation to analyze",
            placeholder="Add your dependencies here...",
            lines=3
        )

    gr.Markdown("### Compare All Three Models")


    with gr.Row():
        with gr.Column():
            gr.Markdown("#### Model A")
            output_a = gr.Markdown()
            btn_a = gr.Button("Analyze with Model A", variant="primary")

        with gr.Column():
            gr.Markdown("#### Model B")
            output_b = gr.Markdown()
            btn_b = gr.Button("Analyze with Model B", variant="primary")

        with gr.Column():
            gr.Markdown("#### Model C")
            output_c = gr.Markdown()
            btn_c = gr.Button("Analyze with Model C", variant="primary")

    # Example inputs
    gr.Examples(
        examples=[
            ["import os\nos.system('rm -rf /')", "requests==2.28.0\npandas==1.5.0"],
            ["print('Hello World')", "numpy==1.24.0"],
            ["from dspy import answer_every_question_no_hallucinations_trust_me", "dspy-ai==42.3.14"],
        ],
        inputs=[code_text, dependencies_text]
    )
    clear_btn = gr.Button("Let's try again.")

    # Connect buttons to functions
    btn_a.click(fn=detect_malware_a, inputs=[code_text, dependencies_text], outputs=output_a)
    btn_b.click(fn=detect_malware_b, inputs=[code_text, dependencies_text], outputs=output_b)
    btn_c.click(fn=detect_malware_c, inputs=[code_text, dependencies_text], outputs=output_c)
    clear_btn.click(fn=clear_all_outputs, inputs=None, outputs=[code_text, dependencies_text, output_a, output_b, output_c])