import gradio as gr
import requests
import os

# --- 1. DYNAMIC API ROUTING ---
# Pulls the live AMD IP from Hugging Face Settings
API_URL = os.environ.get("INTELLIGUARD_API", "http://127.0.0.1:8000/scan")

# Custom CSS
custom_css = """
.gradio-container { background-color: #06090f !important; color: white !important; }
.fake-sidebar { 
    background: #0b0f19 !important; 
    border-right: 1px solid #1f2937 !important; 
    padding: 20px !important;
    min-width: 250px !important;
}
.status-badge { 
    background: rgba(16, 185, 129, 0.1); border: 1px solid #10b981; color: #10b981; 
    padding: 8px; border-radius: 20px; text-align: center; margin-bottom: 30px;
}
.card { 
    background: #111827; padding: 20px; border-radius: 12px; text-align: center;
    border-top: 3px solid #3b82f6; margin-bottom: 10px;
}
.establish-btn { background: linear-gradient(90deg, #3b82f6 0%, #00d4ff 100%) !important; color: white !important; font-weight: bold !important;}
/* Fix for dark mode inputs */
input, textarea { color: white !important; background-color: #161b22 !important; }
"""

# --- 2. THE BRAIN LOGIC ---
def process_query(query, tot_count, thr_count, saf_count, current_log):
    if not query.strip():
        return tot_count, thr_count, saf_count, "", "Please enter a query.", current_log, current_log, gr.update(), gr.update(), gr.update()

    tot_count += 1
    
    try:
        # Send payload to AMD Server
        res = requests.post(API_URL, json={"text": query}, timeout=15)
        data = res.json()
        verdict = data.get("verdict", "ERROR")
        category = data.get("attack_category", "Unknown")
        score = data.get("score", 0.0)
    except Exception as e:
        verdict = "ERROR"
        category = "Connection Failed"
        score = 0.0

    # Logic for metrics and UI updating
    if verdict == "INJECTION":
        thr_count += 1
        color = "#ef4444" # Red
        ui_status = f"🛑 THREAT BLOCKED: {category} (Confidence: {score:.2f})"
        new_log = f"<div style='border-left: 3px solid #ef4444; padding-left: 10px; margin-bottom: 10px; font-size: 12px;'><b>Threat:</b> {category}<br><span style='color: #9ca3af;'>{query[:30]}...</span></div>"
    elif verdict == "ERROR":
        color = "#eab308" # Yellow
        ui_status = f"⚠️ CONNECTION ERROR: Ensure AMD backend is running and Port 8000 is open."
        new_log = f"<div style='border-left: 3px solid #eab308; padding-left: 10px; margin-bottom: 10px; font-size: 12px;'><b>Error:</b> API Timeout</div>"
    else:
        saf_count += 1
        color = "#10b981" # Green
        ui_status = f"✅ SAFE: Query permitted to internal LLM."
        new_log = f"<div style='border-left: 3px solid #10b981; padding-left: 10px; margin-bottom: 10px; font-size: 12px;'><b>Safe Query</b><br><span style='color: #9ca3af;'>{query[:30]}...</span></div>"

    # Compile the updated HTML components
    updated_log = new_log + current_log
    
    tot_html = f"<div class='card' style='border-color:#3b82f6'><div style='font-size:24px'>{tot_count}</div><div style='font-size:10px; color:#9ca3af'>Total Queries</div></div>"
    thr_html = f"<div class='card' style='border-color:#ef4444'><div style='font-size:24px'>{thr_count}</div><div style='font-size:10px; color:#9ca3af'>Threats Blocked</div></div>"
    saf_html = f"<div class='card' style='border-color:#10b981'><div style='font-size:24px'>{saf_count}</div><div style='font-size:10px; color:#9ca3af'>Safe Queries</div></div>"
    
    out_box = f"<div style='padding: 15px; border-radius: 8px; background: #161b22; border: 1px solid {color}; color: {color}; text-align: center; font-weight: bold;'>{ui_status}</div>"

    # Return order MUST match the outputs list below
    return tot_count, thr_count, saf_count, "", out_box, updated_log, updated_log, tot_html, thr_html, saf_html


# --- 3. THE UI WIRING ---
with gr.Blocks(css=custom_css) as demo:
    # State variables to track numbers
    tot_state = gr.State(0)
    thr_state = gr.State(0)
    saf_state = gr.State(0)
    log_state = gr.State("")

    with gr.Row():
        # --- SIDEBAR ---
        with gr.Column(scale=1, elem_classes="fake-sidebar"):
            gr.HTML("<div style='font-size: 40px; text-align:center;'>🛡️</div>")
            gr.Markdown("## **INTELLIGUARD**\nSECURITY PLATFORM")
            gr.HTML("<div class='status-badge'>● INTELLIGUARD PROTECTED</div>")
            
            gr.Markdown("### 📋 AUDIT LOG")
            audit_log_ui = gr.HTML("<p style='color: #6b7280;'>No activity recorded.</p>")
            
            gr.Markdown("---")
            gr.Markdown("### 📄 DOCUMENT UPLOAD")
            gr.File(label=None)

        # --- MAIN CONTENT ---
        with gr.Column(scale=4):
            gr.HTML("""
                <div style="text-align: center; margin-top: 20px;">
                    <h1 style="font-size: 32px;">TechCorp Employee Portal</h1>
                    <p style="color: #6b7280;">Secured by <span style="color: #10b981;">IntelliGuard</span></p>
                </div>
            """)

            # Dynamic Metric Cards
            with gr.Row():
                tot_ui = gr.HTML("<div class='card' style='border-color:#3b82f6'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Total Queries</div></div>")
                thr_ui = gr.HTML("<div class='card' style='border-color:#ef4444'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Threats Blocked</div></div>")
                saf_ui = gr.HTML("<div class='card' style='border-color:#10b981'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Safe Queries</div></div>")

            with gr.Accordion("📧 Enterprise Email Gateway", open=True):
                email = gr.Textbox(label="Target Email", value="security-ops@techcorp.com", interactive=True)
                connect_btn = gr.Button("🔗 Establish Secure Connection", elem_classes="establish-btn")

            gr.Markdown("<div style='height: 80px;'></div>")
            
            # Action Area
            status_box = gr.HTML("<div style='padding: 15px; text-align: center; color: #6b7280;'>Awaiting Payload...</div>")
            
            with gr.Row():
                chat = gr.Textbox(show_label=False, placeholder="Ask about company policies...", scale=10, interactive=True)
                submit = gr.Button("↑", scale=1, variant="primary")

    # Connect the UI to the Python Function
    submit.click(
        fn=process_query,
        inputs=[chat, tot_state, thr_state, saf_state, log_state],
        outputs=[tot_state, thr_state, saf_state, chat, status_box, log_state, audit_log_ui, tot_ui, thr_ui, saf_ui]
    )
    chat.submit(
        fn=process_query,
        inputs=[chat, tot_state, thr_state, saf_state, log_state],
        outputs=[tot_state, thr_state, saf_state, chat, status_box, log_state, audit_log_ui, tot_ui, thr_ui, saf_ui]
    )

demo.launch()