import gradio as gr
import os
import uuid
from dotenv import load_dotenv

load_dotenv()

from graph import create_reasoning_graph
from state import create_initial_state

graph = create_reasoning_graph()

CYBERPUNK_CSS = """
body { 
    background-color: #0b0f19; 
    color: #e0e0e0; 
}
.header { 
    text-align: center; 
    margin-bottom: 30px; 
    padding: 20px; 
}
.header h1 { 
    color: #6366f1; 
    font-family: 'Courier New', monospace; 
    font-weight: bold; 
    font-size: 2.5em; 
    text-shadow: 0 0 10px rgba(99, 102, 241, 0.5); 
}
.secure-box { 
    border: 1px solid #374151; 
    background-color: #111827; 
    padding: 20px; 
    border-radius: 12px; 
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.3); 
}
.json-box { 
    font-family: 'Courier New', monospace; 
    font-size: 13px; 
    background-color: #000; 
    color: #0f0; 
    border: 1px solid #333; 
}
"""

def run_agent(user_input):
    """Stream agent execution with safe key access."""
    if not user_input or not user_input.strip():
        yield "❌ Please enter a task.", "⚠️ System Ready - Awaiting Input"
        return
    
    # Immediate feedback for responsiveness
    yield "🔄 System Initializing...", "⏳ Initializing Secure Ledger..."
    
    task_id = f"task_{uuid.uuid4().hex[:8]}"
    reasoning_log = ""
    crypto_log = "╔════════════════════════════════════════════════════════╗\n"
    crypto_log += "║           SECURE IMMUTABLE LEDGER INITIALIZED         ║\n"
    crypto_log += "╚════════════════════════════════════════════════════════╝\n\n"
    
    initial_state = create_initial_state(
        task=user_input.strip(),
        task_id=task_id,
        user_id="gradio_user"
    )
    
    try:
        final_state = graph.invoke(initial_state)
        
        # Build Reasoning Output (Markdown)
        reasoning_log = "## 🎯 Execution Plan\n\n"
        
        plan = final_state.get("plan")
        if plan and hasattr(plan, "steps"):
            for step in plan.steps:
                step_num = getattr(step, "step_number", "?")
                action = getattr(step, "action", "Unknown action")
                reasoning_log += f"**Step {step_num}:** {action}\n\n"
        else:
            reasoning_log += "_No plan available._\n\n"
        
        # Reasoning Steps
        reasoning_log += "## 📝 Reasoning Trace\n\n"
        justifications = final_state.get("justifications", [])
        if justifications:
            for j in justifications:
                step_num = getattr(j, "step_number", "?")
                reasoning = getattr(j, "reasoning", "No reasoning provided")
                reasoning_log += f"**Step {step_num}:**\n\n{reasoning}\n\n---\n\n"
        else:
            reasoning_log += "_No reasoning steps available._\n\n"
        
        # Status
        status = final_state.get("status", "unknown")
        if status == "completed":
            reasoning_log += "✅ **Task Completed Successfully**"
        elif status == "blocked":
            reasoning_log += "🚫 **Task Blocked by Safety Guardrails**"
        elif status == "failed":
            error = final_state.get("error", "Unknown error")
            reasoning_log += f"❌ **Task Failed:** {error}"
        else:
            reasoning_log += f"⚠️ **Status:** {status}"
        
        # Build Crypto Ledger Output
        logs = final_state.get("logs", [])
        if logs:
            for idx, log_entry in enumerate(logs, 1):
                timestamp = getattr(log_entry, "timestamp", "N/A")
                if hasattr(timestamp, "isoformat"):
                    timestamp = timestamp.isoformat()
                
                action_hash = str(getattr(log_entry, "action_hash", "N/A"))
                merkle_root = str(getattr(log_entry, "merkle_root", "N/A"))
                worm_path = getattr(log_entry, "worm_path", "memory")
                step_number = getattr(log_entry, "step_number", "?")
                
                crypto_log += f"╔════════════════════════════════════════════════════════╗\n"
                crypto_log += f"║ BLOCK #{idx}                                                  ║\n"
                crypto_log += f"╠════════════════════════════════════════════════════════╣\n"
                crypto_log += f"║ STEP:      {str(step_number):<48}║\n"
                crypto_log += f"║ TIMESTAMP: {str(timestamp):<48}║\n"
                crypto_log += f"║ HASH:      {action_hash[:48]:<48}║\n"
                crypto_log += f"║ ROOT:      {merkle_root[:48]:<48}║\n"
                crypto_log += f"║ WORM:      {str(worm_path):<48}║\n"
                crypto_log += f"║ STATUS:    VERIFIED ✓                                    ║\n"
                crypto_log += f"╚════════════════════════════════════════════════════════╝\n\n"
        else:
            crypto_log += "⏳ No cryptographic logs generated yet.\n"
        
        yield reasoning_log, crypto_log
        
    except Exception as e:
        error_msg = f"❌ **Execution Error:** {str(e)}"
        crypto_log += f"\n╔════════════════════════════════════════════════════════╗\n"
        crypto_log += f"║ ERROR: {str(e)[:48]:<48}║\n"
        crypto_log += f"╚════════════════════════════════════════════════════════╝\n"
        yield error_msg, crypto_log


with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate"),
    css=CYBERPUNK_CSS,
    title="Secure Reasoning MCP"
) as demo:
    
    # Header
    gr.HTML("<div class='header'><h1>🛡️ Secure Reasoning MCP Server</h1></div>")
    
    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown("### 📋 Task Input")
            user_input = gr.Textbox(
                label="Task Description",
                placeholder="E.g.: Analyze the benefits of renewable energy...",
                lines=4,
                elem_classes="secure-box"
            )
            submit_btn = gr.Button(
                "▶️ Start Secure Reasoning",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("### 🎯 Agent Reasoning Flow")
            reasoning_output = gr.Markdown(
                value="⏳ Waiting for task...",
                elem_classes="secure-box"
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### 🔐 Immutable Crypto Ledger")
            crypto_output = gr.Textbox(
                value="⏳ System Ready - Awaiting task execution...",
                lines=24,
                max_lines=30,
                show_label=False,
                interactive=False,
                elem_classes="json-box"
            )
    
    gr.Examples(
        examples=[
            ["Analyze the top 3 benefits of AI in healthcare"],
            ["Create a simple Python function to calculate fibonacci numbers"],
            ["Explain quantum computing in simple terms"]
        ],
        inputs=user_input
    )
    
    submit_btn.click(
        fn=run_agent,
        inputs=[user_input],
        outputs=[reasoning_output, crypto_output]
    )

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