app.py

import os
import asyncio
from textwrap import dedent
from dotenv import load_dotenv
from crewai import Agent, Crew, Process, Task
from langchain_openai import ChatOpenAI
import gradio as gr

# Load the .env file to get the API key
load_dotenv()

# Verify API key exists
if not os.getenv("OPENAI_API_KEY"):
    raise ValueError("OPENAI_API_KEY not found in environment variables. Please add it to your .env file.")

# Initialize LLM
llm = ChatOpenAI(
    model="gpt-4-turbo-preview",
    temperature=0.1,
    api_key=os.getenv("OPENAI_API_KEY")
)

# --- 1. Agent Definitions ---
lead_architect = Agent(
    role="Lead Network Architect",
    goal=dedent("""
        Translate high-level business and application needs into a detailed list
        of specific, actionable technical requirements for a new data center network.
    """),
    backstory=dedent("""
        You are a seasoned Lead Architect who excels at interfacing with business stakeholders.
        Your primary skill is listening to business goals and application needs, then distilling
        them into a clear and comprehensive list of technical specifications that engineers can build from.
        You think in terms of availability, security, scalability, and performance metrics.
    """),
    llm=llm,
    allow_delegation=False,
    verbose=True,
)

senior_network_architect = Agent(
    role="Senior Network Architect", 
    goal=dedent("""
        Take a detailed list of technical requirements and create three distinct, high-level
        data center network designs. Each design must be for a different major vendor
        (Cisco, Arista, Juniper) and must meet all specified technical requirements.
    """),
    backstory=dedent("""
        You are a hands-on Senior Network Architect with deep product knowledge across
        Cisco, Arista, and Juniper. You are an expert in creating detailed High-Level Designs (HLDs)
        based on a given set of technical specifications. You focus purely on the technical implementation.
    """),
    llm=llm,
    allow_delegation=False,
    verbose=True,
)

peer_review_architect = Agent(
    role="Peer Review Architect",
    goal=dedent("""
        Critically review three network designs against the original technical requirements
        to validate their feasibility, compare their trade-offs, and select the optimal solution.
    """),
    backstory=dedent("""
        You are an impartial and meticulous architect responsible for quality control.
        Your job is to act as a peer reviewer. You take a set of designs and the original
        requirements, and your task is to find any gaps, inconsistencies, or trade-offs.
        You provide the final validation and reasoned recommendation on which design to proceed with.
    """),
    llm=llm,
    allow_delegation=False,
    verbose=True,
)

# --- 2. Task Definitions ---
def create_tasks(business_requirements, application_requirements):
    """Create tasks with proper input context"""
    
    initial_requirements = f"Business Requirements: {business_requirements}\nApplication Requirements: {application_requirements}"
    
    technical_translation_task = Task(
        description=dedent(f"""
            Analyze the following business and application requirements and translate
            them into a detailed list of technical network requirements.

            Create a comprehensive table with three columns: 'Category', 'Requirement', and 'Justification'.
            
            Categories should include: Availability, Performance, Security, Scalability, Management, and Budget.
            
            Requirements to analyze:
            {initial_requirements}
            
            Focus on creating specific, measurable technical requirements that can guide network design decisions.
        """),
        expected_output="A detailed markdown table with three columns: 'Category', 'Requirement', and 'Justification'.",
        agent=lead_architect,
    )

    design_creation_task = Task(
        description=dedent("""
            Based on the technical requirements from the Lead Architect, create three distinct 
            high-level data center designs using a spine-leaf architecture.

            **CRITICAL:** Prioritize the original business requirements when making design trade-offs.
            If 'Strict Budget' was selected, choose cost-effective hardware even if it means 
            compromising on supporting requirements.
            
            Create one design for each vendor:
            - Cisco design with specific switch models and NX-OS
            - Juniper design with specific switch models and Junos
            - Arista design with specific switch models and EOS
            
            For each design, specify:
            1. Recommended spine switch models and quantities
            2. Recommended leaf switch models and quantities  
            3. Network operating system version
            4. Brief rationale explaining how it meets the technical requirements
            5. Estimated cost range if budget was a consideration
        """),
        expected_output="Three complete high-level designs, one for each vendor (Cisco, Juniper, Arista).",
        agent=senior_network_architect,
        context=[technical_translation_task]
    )

    validation_task = Task(
        description=dedent("""
            Review the technical requirements and the three vendor designs (Cisco, Juniper, Arista) 
            provided by previous team members.

            Create a comprehensive validation report with these sections:

            1. **Requirements Validation**: Check if each vendor design meets all technical requirements
            2. **Comparative Analysis**: Compare pros/cons of Cisco vs Juniper vs Arista solutions
            3. **Trade-off Analysis**: Identify any compromises made in each design
            4. **Final Recommendation**: Select the best vendor solution with clear justification
            5. **Implementation Considerations**: Next steps and potential risks

            Base your analysis solely on the specific vendor designs provided, not generic concepts.
        """),
        expected_output="A comprehensive validation report with requirements check, vendor comparison, and final recommendation.",
        agent=peer_review_architect,
        context=[technical_translation_task, design_creation_task]
    )
    
    return [technical_translation_task, design_creation_task, validation_task]

# --- 3. CrewAI Execution Function ---
def run_design_crew(business_reqs_list, app_reqs_list, progress=gr.Progress()):
    """Run the CrewAI design crew with proper execution flow"""
    
    # Input validation
    if not business_reqs_list and not app_reqs_list:
        return (
            "❌ **Error**: Please select at least one requirement.",
            "",
            ""
        )
    
    # Verify API key
    if not os.getenv("OPENAI_API_KEY"):
        return (
            "❌ **Error**: OPENAI_API_KEY not found. Please add it to your .env file.",
            "",
            ""
        )
    
    try:
        # Prepare requirements strings
        business_reqs_str = ", ".join(business_reqs_list) if business_reqs_list else "None specified"
        app_reqs_str = ", ".join(app_reqs_list) if app_reqs_list else "None specified"
        
        progress(0.1, desc="Initializing AI agents...")
        
        # Create tasks
        tasks = create_tasks(business_reqs_str, app_reqs_str)
        
        # Create crew
        crew = Crew(
            agents=[lead_architect, senior_network_architect, peer_review_architect],
            tasks=tasks,
            process=Process.sequential,
            verbose=True,
            memory=True,  # Enable memory for better context retention
        )
        
        progress(0.2, desc="Starting design process...")
        
        # Prepare inputs
        inputs = {
            'business_requirements': business_reqs_str,
            'application_requirements': app_reqs_str,
            'requirements_context': f"Business: {business_reqs_str}\nApplications: {app_reqs_str}"
        }
        
        progress(0.3, desc="Lead Architect analyzing requirements...")
        
        # Execute crew
        result = crew.kickoff(inputs=inputs)
        
        progress(0.9, desc="Finalizing reports...")
        
        # Extract individual task outputs
        task_outputs = result.tasks_output if hasattr(result, 'tasks_output') else []
        
        if len(task_outputs) >= 3:
            lead_output = task_outputs[0].raw if hasattr(task_outputs[0], 'raw') else str(task_outputs[0])
            senior_output = task_outputs[1].raw if hasattr(task_outputs[1], 'raw') else str(task_outputs[1])
            peer_output = task_outputs[2].raw if hasattr(task_outputs[2], 'raw') else str(task_outputs[2])
        else:
            # Fallback if task outputs aren't properly separated
            full_output = str(result)
            # Try to split the output into sections
            sections = full_output.split("# Task")
            if len(sections) >= 3:
                lead_output = f"# Technical Requirements\n{sections[1]}" if len(sections) > 1 else "Output processing..."
                senior_output = f"# Network Designs\n{sections[2]}" if len(sections) > 2 else "Processing designs..."
                peer_output = f"# Validation Report\n{sections[3]}" if len(sections) > 3 else "Finalizing validation..."
            else:
                lead_output = "✅ Requirements analysis complete"
                senior_output = "✅ Designs generated"  
                peer_output = full_output
        
        progress(1.0, desc="Complete!")
        
        return lead_output, senior_output, peer_output
        
    except Exception as e:
        error_msg = f"❌ **Error occurred**: {str(e)}\n\n**Troubleshooting:**\n- Check your OpenAI API key\n- Ensure stable internet connection\n- Try selecting fewer requirements if the request is too complex"
        return error_msg, error_msg, error_msg

# --- 4. Enhanced CSS ---
custom_css = """
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
    
    /* Hide Gradio footer */
    footer { display: none !important; }
    
    /* Main container styling */
    .gradio-container {
        max-width: 1200px !important;
        margin: 0 auto !important;
        padding: 2rem !important;
        background: #ffffff !important;
        font-family: 'Inter', system-ui, -apple-system, sans-serif !important;
        min-height: 100vh !important;
    }
    
    /* Clean container backgrounds */
    .gradio-blocks, .gradio-blocks > div, .gradio-group, .gradio-column, .gradio-row {
        background: transparent !important;
        border: none !important;
        box-shadow: none !important;
    }
    
    /* Header styling */
    .main-header {
        text-align: left !important;
        margin-bottom: 3rem !important;
        background: transparent !important;
    }
    
    .header-title {
        display: flex !important;
        align-items: center !important;
        gap: 0.75rem !important;
        font-size: 1.75rem !important;
        font-weight: 600 !important;
        color: #1f2937 !important;
        margin-bottom: 1rem !important;
        font-family: 'Inter', sans-serif !important;
    }
    
    .bot-icon {
        width: 1.75rem !important;
        height: 1.75rem !important;
        color: #3b82f6 !important;
    }
    
    .header-subtitle {
        color: #6b7280 !important;
        font-size: 1rem !important;
        margin: 0 !important;
        font-weight: 400 !important;
        font-family: 'Inter', sans-serif !important;
        line-height: 1.5 !important;
    }
    
    /* Section headers */
    .section-title {
        color: #3b82f6 !important;
        font-size: 1.125rem !important;
        font-weight: 700 !important;
        margin-bottom: 1.5rem !important;
        margin-top: 2rem !important;
        padding: 0 0 0.5rem 0 !important;
        background: transparent !important;
        font-family: 'Inter', sans-serif !important;
        border-bottom: 2px solid #e5e7eb !important;
    }
    
    .section-title:first-of-type {
        margin-top: 0 !important;
    }
    
    /* Checkbox styling */
    .gradio-checkboxgroup {
        gap: 0.5rem !important;
        background: transparent !important;
        border: none !important;
        padding: 0 !important;
        margin-bottom: 2rem !important;
    }
    
    .gradio-checkboxgroup label {
        background: #f8fafc !important;
        border: 1px solid #e2e8f0 !important;
        border-radius: 0.5rem !important;
        padding: 0.75rem 1rem !important;
        margin: 0 0 0.5rem 0 !important;
        transition: all 0.2s ease !important;
        cursor: pointer !important;
        display: flex !important;
        align-items: center !important;
        gap: 0.75rem !important;
        font-size: 0.9rem !important;
        line-height: 1.4 !important;
        color: #475569 !important;
        font-weight: 400 !important;
        font-family: 'Inter', sans-serif !important;
    }
    
    .gradio-checkboxgroup label:hover {
        background: #f1f5f9 !important;
        border-color: #10b981 !important;
        color: #1e293b !important;
    }
    
    .gradio-checkboxgroup input[type="checkbox"] {
        accent-color: #10b981 !important;
        width: 1rem !important;
        height: 1rem !important;
        margin: 0 !important;
        flex-shrink: 0 !important;
        border-radius: 0.25rem !important;
    }
    
    .gradio-checkboxgroup label:has(input[type="checkbox"]:checked) {
        background: #ecfdf5 !important;
        border-color: #10b981 !important;
        color: #065f46 !important;
        font-weight: 500 !important;
    }
    
    /* Generate button */
    .generate-btn {
        background: linear-gradient(135deg, #10b981 0%, #059669 100%) !important;
        color: white !important;
        border: none !important;
        padding: 1rem 1.5rem !important;
        border-radius: 0.75rem !important;
        font-weight: 600 !important;
        font-size: 1rem !important;
        width: 100% !important;
        margin-top: 2rem !important;
        cursor: pointer !important;
        transition: all 0.3s ease !important;
        font-family: 'Inter', sans-serif !important;
        min-height: 3rem !important;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1) !important;
    }
    
    .generate-btn:hover {
        background: linear-gradient(135deg, #059669 0%, #047857 100%) !important;
        transform: translateY(-2px) !important;
        box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1) !important;
    }
    
    /* Report section titles */
    .report-title {
        color: #1f2937 !important;
        font-size: 1.125rem !important;
        font-weight: 600 !important;
        margin-bottom: 1rem !important;
        margin-top: 2rem !important;
        background: transparent !important;
        font-family: 'Inter', sans-serif !important;
        padding: 0.5rem 0 !important;
        border-bottom: 1px solid #e5e7eb !important;
    }
    
    .report-title:first-of-type {
        margin-top: 0 !important;
    }
    
    /* Markdown output */
    .gradio-markdown {
        background: transparent !important;
        border: none !important;
        padding: 1rem !important;
        font-family: 'Inter', sans-serif !important;
        line-height: 1.6 !important;
        background: #fafafa !important;
        border-radius: 0.5rem !important;
        border: 1px solid #e5e7eb !important;
    }
    
    .gradio-markdown h1, .gradio-markdown h2, .gradio-markdown h3 {
        color: #1f2937 !important;
        font-weight: 600 !important;
        margin-top: 1.5rem !important;
        margin-bottom: 0.5rem !important;
    }
    
    .gradio-markdown table {
        width: 100% !important;
        border-collapse: collapse !important;
        margin: 1rem 0 !important;
        background: white !important;
        border-radius: 0.5rem !important;
        overflow: hidden !important;
        box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1) !important;
    }
    
    .gradio-markdown table th,
    .gradio-markdown table td {
        padding: 0.75rem !important;
        border: 1px solid #e5e7eb !important;
        text-align: left !important;
        vertical-align: top !important;
    }
    
    .gradio-markdown table th {
        background: #f8fafc !important;
        font-weight: 600 !important;
        color: #374151 !important;
    }
    
    .gradio-markdown table tr:nth-child(even) {
        background: #f9fafb !important;
    }
    
    /* Placeholder content */
    .placeholder-content {
        text-align: center !important;
        color: #9ca3af !important;
        padding: 4rem 2rem !important;
        background: #f8fafc !important;
        border: 2px dashed #e5e7eb !important;
        border-radius: 0.75rem !important;
        font-family: 'Inter', sans-serif !important;
    }
    
    .placeholder-icon {
        width: 3rem !important;
        height: 3rem !important;
        margin: 0 auto 1rem auto !important;
        opacity: 0.5 !important;
        color: #9ca3af !important;
    }
    
    .placeholder-content p {
        font-size: 0.9rem !important;
        color: #6b7280 !important;
        margin: 0 !important;
        font-weight: 400 !important;
    }
    
    /* Layout */
    .gradio-row {
        gap: 3rem !important;
        align-items: flex-start !important;
    }
    
    /* Progress styling */
    .gradio-progress {
        background: #f3f4f6 !important;
        border-radius: 0.5rem !important;
        overflow: hidden !important;
    }
    
    .gradio-progress .progress-bar {
        background: linear-gradient(90deg, #10b981, #059669) !important;
        transition: width 0.3s ease !important;
    }
</style>
"""

# --- 5. Gradio Interface ---
with gr.Blocks(
    theme="soft",
    css=custom_css,
    title="AI Co-Designer for Data Center Networks"
) as demo:
    
    # Header
    gr.HTML("""
        <div class="main-header">
            <div class="header-title">
                <svg class="bot-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                    <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9.75 17L9 20l-1 1h8l-1-1-.75-3M3 13h18M5 17h14a2 2 0 002-2V5a2 2 0 00-2-2H5a2 2 0 00-2 2v10a2 2 0 002 2z"/>
                </svg>
                AI Co-Designer for Data Center Networks
            </div>
            <p class="header-subtitle">
                Select your business and application requirements to generate comprehensive network designs with AI-powered validation.
            </p>
        </div>
    """)
    
    # Main layout
    with gr.Row():
        # Left column - Requirements
        with gr.Column(scale=1):
            # Business Requirements
            gr.HTML('<div class="section-title">Business Requirements</div>')
            business_checkboxes = gr.CheckboxGroup(
                choices=[
                    "Strict Budget (CAPEX)",
                    "High Scalability", 
                    "High Automation (OPEX)",
                    "High Availability",
                    "Regulatory Compliance"
                ],
                value=[],
                container=False,
                show_label=False
            )
            
            # Application Requirements  
            gr.HTML('<div class="section-title">Application Requirements</div>')
            application_checkboxes = gr.CheckboxGroup(
                choices=[
                    "Large-Scale Virtualization",
                    "AI/ML Workloads", 
                    "Low-Latency / Real-time",
                    "Big Data Clusters",
                    "High East-West Traffic"
                ],
                value=[],
                container=False,
                show_label=False
            )
            
            # Generate Button
            submit_button = gr.Button(
                "🚀 Generate & Validate Design",
                elem_classes=["generate-btn"]
            )
        
        # Right column - Agent outputs
        with gr.Column(scale=2):
            # Lead Architect Output
            gr.HTML('<div class="report-title">💡 Lead Architect: Technical Requirements</div>')
            lead_architect_output = gr.Markdown(
                value="""<div class="placeholder-content">
                    <svg class="placeholder-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                        <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9 12h6m-6 4h6m2 5H7a2 2 0 01-2-2V5a2 2 0 012-2h5.586a1 1 0 01.707.293l5.414 5.414a1 1 0 01.293.707V19a2 2 0 01-2 2z"/>
                    </svg>
                    <p>Requirements analysis will appear here...</p>
                </div>""",
                container=False,
                show_label=False
            )
            
            # Senior Architect Output
            gr.HTML('<div class="report-title">🏗️ Senior Architect: Network Designs</div>')
            senior_architect_output = gr.Markdown(
                value="""<div class="placeholder-content">
                    <svg class="placeholder-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                        <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M19 21V5a2 2 0 00-2-2H7a2 2 0 00-2 2v16m14 0h2m-2 0h-5m-9 0H3m2 0h5M9 7h1m-1 4h1m4-4h1m-1 4h1m-5 10v-5a1 1 0 011-1h2a1 1 0 011 1v5m-4 0h4"/>
                    </svg>
                    <p>Network designs will appear here...</p>
                </div>""",
                container=False,
                show_label=False
            )
            
            # Peer Review Output
            gr.HTML('<div class="report-title">✅ Peer Review: Validation & Recommendation</div>')
            peer_review_output = gr.Markdown(
                value="""<div class="placeholder-content">
                    <svg class="placeholder-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                        <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9 12l2 2 4-4m6 2a9 9 0 11-18 0 9 9 0 0118 0z"/>
                    </svg>
                    <p>Validation report will appear here...</p>
                </div>""",
                container=False,
                show_label=False
            )
    
    # Connect the functionality
    submit_button.click(
        fn=run_design_crew,
        inputs=[business_checkboxes, application_checkboxes],
        outputs=[lead_architect_output, senior_architect_output, peer_review_output],
        show_progress=True
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0", 
        server_port=7860,
        show_error=True
    )