import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

SCAN_TYPES = [
    "Smart Contract Audit",
    "Web Application Security",
    "API Security",
    "Cloud Infrastructure",
    "IoT / OT Device",
    "SAF-T UA Compliance Check",
]

def run_audit(target_description, scan_type, detail_level):
    system_prompt = (
        "You are an AI-powered security auditor for the Audityzer platform. "
        "Perform a structured security audit and return: "
        "1) Executive Summary, 2) Critical Findings (CVSS scored), "
        "3) Medium Findings, 4) Recommendations, 5) Compliance Status (ISO 27001 / DSTU)."
    )
    verbosity = {1: "brief", 2: "standard", 3: "comprehensive"}[detail_level]
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Scan Type: {scan_type}\nDetail Level: {verbosity}\nTarget: {target_description}"},
    ]
    response = ""
    for chunk in client.chat_completion(messages, max_tokens=768, stream=True):
        token = chunk.choices[0].delta.content
        if token:
            response += token
    return response

with gr.Blocks(title="Audityzer Demo", theme=gr.themes.Monochrome()) as demo:
    gr.Markdown(
        """# 🔍 Audityzer — AI Security Audit Platform
        **Automated vulnerability scanning & compliance auditing**  
        Part of AuditorSEC ecosystem | [auditorsec.com](https://auditorsec.com)
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            target = gr.Textbox(
                label="Target / System Description",
                placeholder="Describe the system, contract, or API endpoint to audit...",
                lines=5
            )
            scan_type = gr.Dropdown(
                choices=SCAN_TYPES,
                label="Scan Type",
                value="Smart Contract Audit"
            )
            detail = gr.Radio(
                choices=[1, 2, 3],
                label="Detail Level (1=Brief, 2=Standard, 3=Comprehensive)",
                value=2
            )
            audit_btn = gr.Button("🚀 Run Audit", variant="primary")
        with gr.Column(scale=2):
            result = gr.Textbox(label="Audit Report", lines=20)
    audit_btn.click(run_audit, inputs=[target, scan_type, detail], outputs=result)

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