import gradio as gr
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from datasets import load_dataset
import json

# Global variables to cache data
cached_data = {}
current_language = "EN"

def load_data_for_language(language):
    """Load data from HuggingFace datasets for the specified language."""
    cache_key = f"data_{language}"

    if cache_key in cached_data:
        return cached_data[cache_key]

    dataset_name = "AYI-NEDJIMI/owasp-top10-en" if language == "EN" else "AYI-NEDJIMI/owasp-top10-fr"

    try:
        vulnerabilities = load_dataset(dataset_name, data_files="vulnerabilities.json", split="train")
        cwe_mapping = load_dataset(dataset_name, data_files="cwe_mapping.json", split="train")
        attack_scenarios = load_dataset(dataset_name, data_files="attack_scenarios.json", split="train")
        qa_dataset = load_dataset(dataset_name, data_files="qa_dataset.json", split="train")

        # Convert to pandas DataFrames
        vuln_df = vulnerabilities.to_pandas()
        cwe_df = cwe_mapping.to_pandas()
        attack_df = attack_scenarios.to_pandas()
        qa_df = qa_dataset.to_pandas()

        # Convert list fields to strings for display
        for col in ['cwe_ids', 'keywords', 'steps']:
            if col in vuln_df.columns:
                vuln_df[col] = vuln_df[col].apply(lambda x: ', '.join(x) if isinstance(x, list) else str(x))
            if col in attack_df.columns:
                attack_df[col] = attack_df[col].apply(lambda x: ', '.join(x) if isinstance(x, list) else str(x))
            if col in qa_df.columns:
                qa_df[col] = qa_df[col].apply(lambda x: ', '.join(x) if isinstance(x, list) else str(x))

        data = {
            'vulnerabilities': vuln_df,
            'cwe_mapping': cwe_df,
            'attack_scenarios': attack_df,
            'qa_dataset': qa_df
        }

        cached_data[cache_key] = data
        return data

    except Exception as e:
        print(f"Error loading data for {language}: {e}")
        return None


def get_unique_categories(df, language):
    """Get unique OWASP categories from dataframe."""
    if df is None or 'category' not in df.columns:
        return []
    return sorted(df['category'].unique().tolist())


def get_unique_severities(df, language):
    """Get unique severity levels from dataframe."""
    if df is None or 'severity' not in df.columns:
        return []
    return sorted(df['severity'].unique().tolist())


def filter_dataframe(df, search_text, category_filter, severity_filter):
    """Apply search and filter to dataframe."""
    if df is None:
        return pd.DataFrame()

    result_df = df.copy()

    # Text search across all columns
    if search_text and search_text.strip():
        search_lower = search_text.lower()
        mask = result_df.astype(str).apply(lambda x: x.str.contains(search_lower, case=False)).any(axis=1)
        result_df = result_df[mask]

    # Category filter
    if category_filter and category_filter != "All":
        if 'category' in result_df.columns:
            result_df = result_df[result_df['category'] == category_filter]

    # Severity filter
    if severity_filter and severity_filter != "All":
        if 'severity' in result_df.columns:
            result_df = result_df[result_df['severity'] == severity_filter]

    return result_df


def update_language(language):
    """Update all data when language changes."""
    global current_language
    current_language = language
    data = load_data_for_language(language)
    if data is None:
        return None
    return data


def create_vulnerability_tab(data):
    """Create the Vulnerabilities tab content."""
    if data is None or 'vulnerabilities' not in data:
        return gr.Dataframe(value=pd.DataFrame())

    vuln_df = data['vulnerabilities']
    categories = get_unique_categories(vuln_df, current_language)
    severities = get_unique_severities(vuln_df, current_language)

    with gr.Group():
        gr.Markdown("### Search and Filters")
        with gr.Row():
            search_input = gr.Textbox(label="Search", placeholder="Search vulnerabilities...")
            category_dropdown = gr.Dropdown(["All"] + categories, value="All", label="OWASP Category")

        with gr.Row():
            severity_dropdown = gr.Dropdown(["All"] + severities, value="All", label="Severity")

        dataframe = gr.Dataframe(
            value=vuln_df,
            interactive=False,
            label="Vulnerabilities"
        )

        def update_vuln_table(search, category, severity):
            filtered = filter_dataframe(vuln_df, search, category, severity)
            return filtered

        search_input.change(
            update_vuln_table,
            inputs=[search_input, category_dropdown, severity_dropdown],
            outputs=dataframe
        )
        category_dropdown.change(
            update_vuln_table,
            inputs=[search_input, category_dropdown, severity_dropdown],
            outputs=dataframe
        )
        severity_dropdown.change(
            update_vuln_table,
            inputs=[search_input, category_dropdown, severity_dropdown],
            outputs=dataframe
        )

    return dataframe


def create_cwe_mapping_tab(data):
    """Create the CWE Mapping tab content."""
    if data is None or 'cwe_mapping' not in data:
        return gr.Dataframe(value=pd.DataFrame())

    cwe_df = data['cwe_mapping']
    categories = get_unique_categories(cwe_df, current_language)

    with gr.Group():
        gr.Markdown("### CWE Mapping Explorer")
        with gr.Row():
            search_input = gr.Textbox(label="Search", placeholder="Search CWEs...")
            category_dropdown = gr.Dropdown(["All"] + categories, value="All", label="OWASP Category")

        dataframe = gr.Dataframe(
            value=cwe_df,
            interactive=False,
            label="CWE Mappings"
        )

        def update_cwe_table(search, category):
            filtered = filter_dataframe(cwe_df, search, category, None)
            return filtered

        search_input.change(
            update_cwe_table,
            inputs=[search_input, category_dropdown],
            outputs=dataframe
        )
        category_dropdown.change(
            update_cwe_table,
            inputs=[search_input, category_dropdown],
            outputs=dataframe
        )

    return dataframe


def create_attack_scenarios_tab(data):
    """Create the Attack Scenarios tab content."""
    if data is None or 'attack_scenarios' not in data:
        return gr.Dataframe(value=pd.DataFrame())

    attack_df = data['attack_scenarios']
    categories = get_unique_categories(attack_df, current_language)

    with gr.Group():
        gr.Markdown("### Attack Scenarios")
        with gr.Row():
            search_input = gr.Textbox(label="Search", placeholder="Search scenarios...")
            category_dropdown = gr.Dropdown(["All"] + categories, value="All", label="OWASP Category")

        dataframe = gr.Dataframe(
            value=attack_df,
            interactive=False,
            label="Attack Scenarios"
        )

        def update_attack_table(search, category):
            filtered = filter_dataframe(attack_df, search, category, None)
            return filtered

        search_input.change(
            update_attack_table,
            inputs=[search_input, category_dropdown],
            outputs=dataframe
        )
        category_dropdown.change(
            update_attack_table,
            inputs=[search_input, category_dropdown],
            outputs=dataframe
        )

    return dataframe


def create_qa_tab(data):
    """Create the Q&A tab content."""
    if data is None or 'qa_dataset' not in data:
        return gr.Dataframe(value=pd.DataFrame())

    qa_df = data['qa_dataset']
    categories = get_unique_categories(qa_df, current_language)

    with gr.Group():
        gr.Markdown("### Q&A Dataset")
        with gr.Row():
            search_input = gr.Textbox(label="Search", placeholder="Search Q&A...")
            category_dropdown = gr.Dropdown(["All"] + categories, value="All", label="Category")

        dataframe = gr.Dataframe(
            value=qa_df,
            interactive=False,
            label="Q&A"
        )

        def update_qa_table(search, category):
            filtered = filter_dataframe(qa_df, search, category, None)
            return filtered

        search_input.change(
            update_qa_table,
            inputs=[search_input, category_dropdown],
            outputs=dataframe
        )
        category_dropdown.change(
            update_qa_table,
            inputs=[search_input, category_dropdown],
            outputs=dataframe
        )

    return dataframe


def create_statistics_tab(data):
    """Create the Statistics tab with visualizations."""
    if data is None:
        return gr.Plot(value=None)

    vuln_df = data['vulnerabilities']

    with gr.Group():
        gr.Markdown("### OWASP Top 10 Statistics")

        # Create chart tabs
        with gr.Tabs():
            with gr.TabItem("CWEs per Category"):
                fig = create_cwe_chart(vuln_df)
                gr.Plot(value=fig)

            with gr.TabItem("Severity Distribution"):
                fig = create_severity_chart(vuln_df)
                gr.Plot(value=fig)

            with gr.TabItem("Vulnerabilities per Category"):
                fig = create_category_chart(vuln_df)
                gr.Plot(value=fig)


def create_cwe_chart(df):
    """Create CWE distribution chart."""
    if 'category' not in df.columns or 'cwe_ids' not in df.columns:
        return go.Figure()

    cwe_count = df.groupby('category').size().sort_values(ascending=False)
    fig = go.Figure(data=[go.Bar(x=cwe_count.index, y=cwe_count.values)])
    fig.update_layout(
        title="CWEs per OWASP Category",
        xaxis_title="OWASP Category",
        yaxis_title="Count",
        hovermode='x unified'
    )
    return fig


def create_severity_chart(df):
    """Create severity distribution chart."""
    if 'severity' not in df.columns:
        return go.Figure()

    severity_count = df['severity'].value_counts().sort_values(ascending=False)
    fig = go.Figure(data=[
        go.Bar(
            x=severity_count.index,
            y=severity_count.values,
            marker_color=['red', 'orange', 'yellow', 'green'][:len(severity_count)]
        )
    ])
    fig.update_layout(
        title="Severity Distribution",
        xaxis_title="Severity Level",
        yaxis_title="Count",
        hovermode='x unified'
    )
    return fig


def create_category_chart(df):
    """Create vulnerabilities per category chart."""
    if 'category' not in df.columns:
        return go.Figure()

    category_count = df['category'].value_counts().sort_values(ascending=False)
    fig = go.Figure(data=[go.Bar(x=category_count.index, y=category_count.values)])
    fig.update_layout(
        title="Vulnerabilities per OWASP Category",
        xaxis_title="OWASP Category",
        yaxis_title="Count",
        hovermode='x unified'
    )
    return fig


def main():
    """Main Gradio application."""
    # Load initial data
    data = load_data_for_language("EN")

    with gr.Blocks(title="OWASP Top 10 Explorer", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # OWASP Top 10 Explorer
        A comprehensive interactive tool for exploring OWASP Top 10 vulnerabilities, CWE mappings, attack scenarios, and security Q&A.
        """)

        # Language selector
        with gr.Row():
            language_radio = gr.Radio(
                choices=["EN", "FR"],
                value="EN",
                label="Language / Langue",
                scale=1
            )

        gr.Markdown("---")

        # Main tabs
        with gr.Tabs():
            with gr.TabItem("Vulnerabilities"):
                vuln_outputs = create_vulnerability_tab(data)

            with gr.TabItem("CWE Mapping"):
                cwe_outputs = create_cwe_mapping_tab(data)

            with gr.TabItem("Attack Scenarios"):
                attack_outputs = create_attack_scenarios_tab(data)

            with gr.TabItem("Q&A"):
                qa_outputs = create_qa_tab(data)

            with gr.TabItem("Statistics"):
                create_statistics_tab(data)

        gr.Markdown("---")

        # Footer
        gr.HTML("""
        <div style='text-align:center; padding:20px; color:#666;'>
            <p>Created by <a href='https://www.ayinedjimi-consultants.fr' target='_blank'>Ayi NEDJIMI</a> - Senior Offensive Cybersecurity & AI Consultant</p>
            <p><a href='https://www.linkedin.com/in/ayi-nedjimi' target='_blank'>LinkedIn</a> | <a href='https://github.com/ayinedjimi' target='_blank'>GitHub</a> | <a href='https://x.com/AyiNEDJIMI' target='_blank'>Twitter/X</a></p>
        </div>
        """)

        # Language change handler
        def on_language_change(language):
            new_data = update_language(language)
            # Return updated dataframes for each tab
            if new_data:
                return (
                    new_data['vulnerabilities'],
                    new_data['cwe_mapping'],
                    new_data['attack_scenarios'],
                    new_data['qa_dataset']
                )
            return (
                pd.DataFrame(),
                pd.DataFrame(),
                pd.DataFrame(),
                pd.DataFrame()
            )

        language_radio.change(
            on_language_change,
            inputs=language_radio,
            outputs=[vuln_outputs, cwe_outputs, attack_outputs, qa_outputs]
        )

    return demo


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