app.py

from __future__ import annotations

import gc
import html
import os
import threading
from functools import lru_cache
from typing import Any

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ============================================================
# Hugging Face Spaces / local compatibility
# ============================================================

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        """Fallback decorator for local execution without ZeroGPU."""

        @staticmethod
        def GPU(function=None, **_kwargs):
            if function is not None:
                return function

            def decorator(func):
                return func

            return decorator

    spaces = _SpacesFallback()


# ============================================================
# Configuration
# ============================================================

MODEL_ID = os.getenv("MODEL_ID", "wasmdashai/asg-v1")
APP_TITLE = "ASGTransformer"
APP_SUBTITLE = "Enterprise Defensive Cybersecurity Scenario Generator"

MODEL_LOCK = threading.Lock()

EXAMPLES = [
    (
        "Create an authorized defensive enterprise scenario focused on "
        "phishing awareness, credential protection, and response readiness."
    ),
    (
        "Design a ransomware tabletop exercise for a financial institution, "
        "including detection, containment, communication, recovery, and lessons learned."
    ),
    (
        "Generate a defensive ICS security scenario involving unusual network "
        "activity, operator validation, segmentation, and incident response coordination."
    ),
    (
        "أنشئ سيناريو دفاعيًا مصرحًا به لمؤسسة مالية يركز على التصيد الاحتيالي، "
        "وحماية بيانات الاعتماد، والكشف المبكر، والاستجابة للحوادث."
    ),
]


# ============================================================
# Theme and CSS
# Gradio 6 requires theme/css to be passed to launch().
# ============================================================

THEME = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="slate",
    neutral_hue="slate",
)

CUSTOM_CSS = """
:root {
    --asg-radius: 18px;
}

.gradio-container {
    max-width: 1500px !important;
    margin: 0 auto !important;
}

#asg-header {
    padding: 30px;
    margin-bottom: 18px;
    border: 1px solid var(--border-color-primary);
    border-radius: 24px;
    background:
        radial-gradient(
            circle at top right,
            rgba(59, 130, 246, 0.18),
            transparent 38%
        ),
        var(--background-fill-primary);
}

#asg-header h1 {
    margin: 0;
    font-size: clamp(2rem, 5vw, 3.5rem);
    letter-spacing: -0.04em;
}

#asg-header p {
    margin: 10px 0 0;
    font-size: 1.05rem;
    opacity: 0.82;
}

.brand-badge {
    display: inline-flex;
    align-items: center;
    gap: 8px;
    padding: 7px 13px;
    margin-bottom: 15px;
    border: 1px solid var(--border-color-primary);
    border-radius: 999px;
    font-size: 0.84rem;
    font-weight: 700;
}

#asg-chatbot {
    min-height: 590px;
    border-radius: var(--asg-radius);
}

#prompt-box textarea {
    font-size: 1rem;
    line-height: 1.65;
}

.metadata-card {
    display: grid;
    grid-template-columns: repeat(2, minmax(0, 1fr));
    gap: 12px;
    padding: 16px;
    border: 1px solid var(--border-color-primary);
    border-radius: var(--asg-radius);
    background: var(--background-fill-secondary);
}

.metadata-item {
    padding: 13px;
    border: 1px solid var(--border-color-primary);
    border-radius: 13px;
    background: var(--background-fill-primary);
}

.metadata-label {
    display: block;
    margin-bottom: 5px;
    font-size: 0.76rem;
    letter-spacing: 0.05em;
    text-transform: uppercase;
    opacity: 0.66;
}

.metadata-value {
    display: block;
    font-weight: 700;
    word-break: break-word;
}

.welcome-card,
.status-warning,
.status-error,
.security-note {
    padding: 16px;
    border: 1px solid var(--border-color-primary);
    border-radius: 14px;
}

.welcome-card,
.security-note {
    background: var(--background-fill-secondary);
}

.status-warning {
    background: rgba(245, 158, 11, 0.12);
}

.status-error {
    background: rgba(239, 68, 68, 0.12);
}

.security-note {
    font-size: 0.9rem;
    line-height: 1.6;
}

footer {
    display: none !important;
}

@media (max-width: 800px) {
    .metadata-card {
        grid-template-columns: 1fr;
    }

    #asg-header {
        padding: 22px;
    }

    #asg-chatbot {
        min-height: 480px;
    }
}
"""


# ============================================================
# Model loading
# ============================================================

@lru_cache(maxsize=1)
def load_model() -> tuple[Any, Any]:
    """
    Load the tokenizer and model once per running Space container.

    trust_remote_code=True is required because ASGTransformer exposes
    the custom generate_scenario() method.
    """
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        torch_dtype="auto",
        device_map="auto",
        low_cpu_mem_usage=True,
    )

    model.eval()

    if tokenizer.pad_token_id is None:
        if tokenizer.eos_token is not None:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({"pad_token": "[PAD]"})
            model.resize_token_embeddings(len(tokenizer))

    return tokenizer, model


# ============================================================
# Helpers
# ============================================================

def language_code(language: str) -> str:
    return "ar" if language == "العربية" else "en"


def defensive_prompt(message: str, code: str) -> str:
    if code == "ar":
        prefix = (
            "أنشئ سيناريو أمن سيبراني دفاعي ومصرح به فقط. ركز على التوعية، "
            "والحماية، والكشف، والاستجابة، والتعافي، ولا تقدم خطوات تشغيلية "
            "تساعد على مهاجمة أنظمة حقيقية.\n\nطلب المستخدم:\n"
        )
    else:
        prefix = (
            "Generate an authorized defensive cybersecurity scenario only. "
            "Focus on awareness, protection, detection, response, and recovery. "
            "Do not provide operational instructions for attacking real systems."
            "\n\nUser request:\n"
        )

    return prefix + message.strip()


def empty_metadata() -> str:
    return """
<div class="welcome-card">
    Scenario type, estimated duration, language, and model details
    will appear here after generation.
</div>
"""


def metadata_card(
    scenario_type: Any,
    duration: Any,
    selected_language: str,
) -> str:
    safe_type = html.escape(str(scenario_type or "Defensive Scenario"))
    safe_duration = html.escape(str(duration or "Not specified"))
    safe_language = html.escape(selected_language)
    safe_model = html.escape(MODEL_ID)

    return f"""
<div class="metadata-card">
    <div class="metadata-item">
        <span class="metadata-label">Scenario Type</span>
        <span class="metadata-value">{safe_type}</span>
    </div>
    <div class="metadata-item">
        <span class="metadata-label">Estimated Duration</span>
        <span class="metadata-value">{safe_duration} minutes</span>
    </div>
    <div class="metadata-item">
        <span class="metadata-label">Language</span>
        <span class="metadata-value">{safe_language}</span>
    </div>
    <div class="metadata-item">
        <span class="metadata-label">Model</span>
        <span class="metadata-value">{safe_model}</span>
    </div>
</div>
"""


def normalize_history(history: Any) -> list[dict[str, Any]]:
    """
    Gradio 6 Chatbot uses OpenAI-style message dictionaries.
    """
    if not history:
        return []

    normalized: list[dict[str, Any]] = []

    for item in history:
        if isinstance(item, dict) and "role" in item and "content" in item:
            normalized.append(
                {
                    "role": str(item["role"]),
                    "content": item["content"],
                }
            )

    return normalized


def clear_conversation() -> tuple[list, str, str]:
    return [], "", empty_metadata()


# ============================================================
# Generation
# ============================================================

@spaces.GPU(duration=180)
def generate_scenario(
    message: str,
    history: list[dict[str, Any]] | None,
    selected_language: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    do_sample: bool,
) -> tuple[str, list[dict[str, Any]], str]:
    history_items = normalize_history(history)
    user_message = (message or "").strip()

    if not user_message:
        return (
            "",
            history_items,
            """
<div class="status-warning">
    Please enter a defensive scenario request before generating.
</div>
""",
        )

    history_items.append(
        {
            "role": "user",
            "content": user_message,
        }
    )

    code = language_code(selected_language)
    prompt = defensive_prompt(user_message, code)

    try:
        with MODEL_LOCK:
            tokenizer, model = load_model()

            generation_options = {
                "language": code,
                "max_new_tokens": int(max_new_tokens),
                "do_sample": bool(do_sample),
            }

            # Sampling-only arguments should not be passed when sampling is disabled.
            if do_sample:
                generation_options.update(
                    {
                        "temperature": float(temperature),
                        "top_p": float(top_p),
                    }
                )

            with torch.inference_mode():
                result = model.generate_scenario(
                    tokenizer,
                    prompt,
                    **generation_options,
                )

        if isinstance(result, dict):
            generated_text = str(result.get("text") or "").strip()
            duration = result.get("estimated_duration_minutes", "Not specified")
            scenario_type = result.get("scenario_type", "Defensive Scenario")
        else:
            generated_text = str(result).strip()
            duration = "Not specified"
            scenario_type = "Defensive Scenario"

        if not generated_text:
            generated_text = (
                "The model completed the request but returned an empty response."
            )

        history_items.append(
            {
                "role": "assistant",
                "content": generated_text,
            }
        )

        return (
            "",
            history_items,
            metadata_card(
                scenario_type=scenario_type,
                duration=duration,
                selected_language=selected_language,
            ),
        )

    except torch.cuda.OutOfMemoryError:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

        error_text = (
            "GPU memory was insufficient. Reduce the maximum output tokens "
            "and submit the request again."
        )
        history_items.append(
            {
                "role": "assistant",
                "content": error_text,
            }
        )

        return (
            "",
            history_items,
            """
<div class="status-error">
    GPU memory limit reached. Try 256 or fewer output tokens.
</div>
""",
        )

    except Exception as error:
        error_text = (
            "The scenario could not be generated. "
            f"{type(error).__name__}: {error}"
        )

        history_items.append(
            {
                "role": "assistant",
                "content": error_text,
            }
        )

        return (
            "",
            history_items,
            f"""
<div class="status-error">
    Generation failed: {html.escape(str(error))}
</div>
""",
        )


# ============================================================
# Gradio 6 interface
# ============================================================

with gr.Blocks(
    title=f"{APP_TITLE} | Defensive Scenario Generator",
    fill_height=True,
) as demo:
    gr.HTML(
        f"""
<section id="asg-header">
    <div class="brand-badge">ASG • Authorized Defensive AI</div>
    <h1>{APP_TITLE}</h1>
    <p>
        {APP_SUBTITLE}. Generate structured, authorized, and
        response-oriented cybersecurity exercises using
        <strong>{html.escape(MODEL_ID)}</strong>.
    </p>
</section>
"""
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=8, elem_id="chat-panel"):
            # Gradio 6 removed the old type="messages" argument.
            # Message dictionaries are now the supported/default format.
            chatbot = gr.Chatbot(
                value=[],
                height=590,
                elem_id="asg-chatbot",
                placeholder=(
                    "Your generated enterprise defensive scenarios "
                    "will appear here."
                ),
                show_label=False,
            )

            message_box = gr.Textbox(
                value="",
                placeholder=(
                    "Describe the authorized defensive scenario you want "
                    "to generate..."
                ),
                lines=4,
                max_lines=10,
                show_label=False,
                elem_id="prompt-box",
            )

            with gr.Row():
                generate_button = gr.Button(
                    "Generate Scenario",
                    variant="primary",
                    size="lg",
                    scale=4,
                )
                clear_button = gr.Button(
                    "Clear",
                    variant="secondary",
                    size="lg",
                    scale=1,
                )

        with gr.Column(scale=4, elem_id="settings-panel"):
            gr.Markdown("## Generation Settings")

            selected_language = gr.Radio(
                choices=["English", "العربية"],
                value="English",
                label="Output language",
            )

            max_new_tokens = gr.Slider(
                minimum=128,
                maximum=1024,
                value=384,
                step=32,
                label="Maximum output tokens",
            )

            do_sample = gr.Checkbox(
                value=True,
                label="Enable creative sampling",
            )

            temperature = gr.Slider(
                minimum=0.1,
                maximum=1.5,
                value=0.7,
                step=0.05,
                label="Temperature",
            )

            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.9,
                step=0.05,
                label="Top-p",
            )

            gr.Markdown("## Scenario Information")
            metadata_output = gr.HTML(empty_metadata())

            gr.HTML(
                """
<div class="security-note">
    <strong>Defensive-use policy</strong><br>
    ASGTransformer is intended for authorized awareness,
    protection, detection, incident response, recovery,
    tabletop exercises, and security training.
</div>
"""
            )

    gr.Markdown("## Example Requests")
    gr.Examples(
        examples=[[example] for example in EXAMPLES],
        inputs=message_box,
        label="Select an example",
    )

    generation_inputs = [
        message_box,
        chatbot,
        selected_language,
        max_new_tokens,
        temperature,
        top_p,
        do_sample,
    ]

    generation_outputs = [
        message_box,
        chatbot,
        metadata_output,
    ]

    generate_button.click(
        fn=generate_scenario,
        inputs=generation_inputs,
        outputs=generation_outputs,
        show_progress="full",
    )

    message_box.submit(
        fn=generate_scenario,
        inputs=generation_inputs,
        outputs=generation_outputs,
        show_progress="full",
    )

    clear_button.click(
        fn=clear_conversation,
        inputs=[],
        outputs=[
            chatbot,
            message_box,
            metadata_output,
        ],
        queue=False,
    )


# ============================================================
# Launch
# ============================================================

if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=1,
        max_size=20,
    ).launch(
        theme=THEME,
        css=CUSTOM_CSS,
    )