app.py

import gradio as gr
import tempfile

from intelligence_pipeline import (
    analyze_threat,
    translate_and_explain,
)
from safety_classifier import safety_check


def handle_threat_intel(message, files):
    base = analyze_threat(message or "")
    if files:
        file_names = ", ".join(f.name for f in files)
        base += f"\n\nAttached files: {file_names}"
    return base


def handle_translation(message, files):
    base = translate_and_explain(message or "")
    if files:
        base += "\n\nNote: Translation mode currently ignores attached files."
    return base


def handle_marketplace_watch(message, files):
    out = f"[Marketplace Watch] Monitoring: {message or ''}"
    if files:
        out += " | attached logs available"
    return out


def handle_analyst_tools(message, files):
    return "[Analyst Tools] " + (message or "")


MODE_ROUTER = {
    "Threat Intel": handle_threat_intel,
    "Translation": handle_translation,
    "Marketplace Watch": handle_marketplace_watch,
    "Analyst Tools": handle_analyst_tools,
}


def process_message(message, files, history, mode):
    text = message or ""

    safety = safety_check(text)
    if safety["blocked"]:
        response = "⚠️ Safety Block: " + safety["reason"]
        history = history + [(f"{mode}: {text}", response)]
        return history, history

    handler = MODE_ROUTER.get(mode, handle_threat_intel)
    response = handler(text, files)

    history = history + [(f"{mode}: {text}", response)]
    return history, history


def clear_chat():
    return [], []


def download_chat(history):
    if not history:
        content = "No conversation yet.\n"
    else:
        lines = []
        for u, s in history:
            lines.append(f"[User]: {u}")
            lines.append(f"[System]: {s}\n")
        content = "\n".join(lines)

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".txt")
    with open(tmp.name, "w", encoding="utf-8") as f:
        f.write(content)
    return tmp.name


with gr.Blocks(title="Threat Landscape Intelligence Console") as demo:
    gr.Markdown("## 🛰 APJ Threat Intelligence Console")

    chat_state = gr.State([])

    mode = gr.Radio(
        ["Threat Intel", "Translation", "Marketplace Watch", "Analyst Tools"],
        value="Threat Intel",
        label="Mode",
        interactive=True,
    )

    chat = gr.Chatbot(label="Dialogue", height=430)

    with gr.Accordion("Attachments & Utilities", open=False):
        file_input = gr.File(label="Upload files", file_count="multiple")
        with gr.Row(variant="compact"):
            clear_btn = gr.Button("Clear Chat", variant="secondary")
            download_btn = gr.Button("Download Transcript", variant="secondary")
        download_file = gr.File(label="Transcript", interactive=False)

    with gr.Row(variant="compact"):
        user_input = gr.Textbox(
            placeholder="Enter text…",
            label="",
            scale=5
        )
        send_btn = gr.Button("Send", variant="primary", scale=1)

    send_btn.click(
        fn=process_message,
        inputs=[user_input, file_input, chat_state, mode],
        outputs=[chat, chat_state],
    ).then(lambda: "", None, user_input)

    user_input.submit(
        fn=process_message,
        inputs=[user_input, file_input, chat_state, mode],
        outputs=[chat, chat_state],
    ).then(lambda: "", None, user_input)

    clear_btn.click(fn=clear_chat, outputs=[chat, chat_state])
    download_btn.click(fn=download_chat, inputs=chat_state, outputs=download_file)


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