import streamlit as st
import requests
import json
import streamlit.components.v1 as components
import os

# Import templates and icons from separate files
from templates import ARCHITECTURE_TEMPLATES, get_all_providers, get_template
from service_icons import SERVICE_ICONS, get_icon_url, map_service_to_icon

# Page Configuration
st.set_page_config(
    page_title="ArchViz AI Pro - Draw.io Style Icons", 
    page_icon="🏗️", 
    layout="wide"
)

# API Configuration
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")

# UI Styling
st.markdown("""
    <style>
    .stTextArea textarea { font-size: 16px; }
    .main-header { 
        text-align: center; 
        color: #155724; 
        background: linear-gradient(135deg, #d4edda 0%, #c3e6cb 100%);
        padding: 15px; 
        border-radius: 10px; 
        margin-bottom: 20px;
        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
    }
    .feature-badge {
        display: inline-block;
        background: #28a745;
        color: white;
        padding: 3px 10px;
        border-radius: 12px;
        font-size: 11px;
        font-weight: bold;
        margin-left: 10px;
    }
    .icon-badge {
        display: inline-block;
        background: #007bff;
        color: white;
        padding: 3px 10px;
        border-radius: 12px;
        font-size: 11px;
        font-weight: bold;
        margin-left: 5px;
    }
    </style>
    <div class="main-header">
        <h1>🏗️ ArchViz AI Pro <span class="feature-badge">Pro Edition</span><span class="icon-badge">Draw.io Icons</span></h1>
        <p>Enterprise Architecture Visualization with Official Cloud Service Icons</p>
    </div>
""", unsafe_allow_html=True)

# Sidebar Configuration
st.sidebar.title("📚 Cloud Architecture Library")
st.sidebar.markdown("Select from professional architecture templates:")

selected_provider = st.sidebar.selectbox(
    "Cloud Provider:",
    ["None"] + get_all_providers()
)

if selected_provider != "None":
    templates = ARCHITECTURE_TEMPLATES[selected_provider]
    selected_template = st.sidebar.selectbox(
        f"{selected_provider} Templates:",
        list(templates.keys())
    )

    if st.sidebar.button("Load Template", use_container_width=True):
        template_data = templates[selected_template]
        st.session_state['template_text'] = template_data["description"]
        st.session_state['template_services'] = template_data.get("services", [])

    st.sidebar.markdown("---")
    st.sidebar.markdown(f"**Template: {selected_template}**")
    template_info = templates[selected_template]
    st.sidebar.info(template_info["description"][:200] + "...")
    st.sidebar.markdown(f"**Use Case:** {template_info['use_case']}")
    
    if "services" in template_info:
        st.sidebar.markdown(f"**Services:** {', '.join(template_info['services'])}")

st.sidebar.markdown("---")
st.sidebar.markdown("### ✨ Pro Features")
st.sidebar.markdown("""
**🎨 Draw.io-Style Icons:**
- ✅ Official AWS Architecture Icons
- ✅ Official GCP Service Icons  
- ✅ Official Azure Service Icons
- ✅ 100+ Service-Specific Icons

**💼 Enterprise Features:**
- ✅ Edit Generated Diagrams
- ✅ High-Quality PNG Export
- ✅ Fullscreen Presentation Mode
- ✅ Professional Templates

**Perfect for:**
- Solution Architecture
- RFP Responses
- Client Presentations
- Technical Documentation
""")

# Main Layout
col1, col2 = st.columns([1, 1.5])

with col1:
    st.subheader("📝 Architecture Description")
    
    default_text = st.session_state.get('template_text', '')
    
    prompt_text = st.text_area(
        "Describe your cloud architecture:",
        value=default_text,
        height=200,
        placeholder="Example: Mobile app connects to AWS API Gateway, which routes to Lambda functions..."
    )

    st.info("💡 **Pro Tip:** Mention specific service names (Lambda, Cloud Run, Functions App) to get official cloud icons!")

    col_gen, col_edit = st.columns(2)
    with col_gen:
        generate_btn = st.button("✨ Generate Diagram", type="primary", use_container_width=True)
    with col_edit:
        if 'mermaid_code' in st.session_state:
            edit_mode = st.button("✏️ Edit Code", use_container_width=True)
            if edit_mode:
                st.session_state['show_editor'] = not st.session_state.get('show_editor', False)

    # Show editor if requested
    if st.session_state.get('show_editor', False) and 'mermaid_code' in st.session_state:
        st.markdown("---")
        st.subheader("🛠️ Edit Diagram Code")
        edited_code = st.text_area(
            "Modify the Mermaid code:",
            value=st.session_state['mermaid_code'],
            height=250,
            key="editor",
            help="Edit the diagram structure, labels, or styling"
        )
        
        col_apply, col_cancel = st.columns(2)
        with col_apply:
            if st.button("✅ Apply Changes", use_container_width=True):
                st.session_state['mermaid_code'] = edited_code
                st.session_state['show_editor'] = False
                st.rerun()
        with col_cancel:
            if st.button("❌ Cancel", use_container_width=True):
                st.session_state['show_editor'] = False
                st.rerun()

    # Display detected services
    if 'template_services' in st.session_state and st.session_state['template_services']:
        st.markdown("---")
        st.markdown("**📦 Template Services:**")
        services_html = ""
        for service in st.session_state['template_services']:
            provider, icon_url = map_service_to_icon(service)
            if icon_url:
                services_html += f'<img src="{icon_url}" width="24" height="24" style="margin: 2px;" title="{service}"> '
        st.markdown(services_html, unsafe_allow_html=True)

# API Call Logic
if generate_btn and prompt_text:
    if not OPENROUTER_API_KEY:
        st.error("⚠️ OPENROUTER_API_KEY not found in environment variables!")
    else:
        with st.spinner("🎨 Creating your architecture diagram with official cloud icons..."):
            try:
                # Build enhanced prompt with icon information
                icon_guidance = """
IMPORTANT - Use service-specific names for proper icons:
AWS: Lambda, API Gateway, DynamoDB, S3, CloudWatch, EC2, RDS, ECS, SQS, SNS, ALB, CloudFront
GCP: Cloud Functions, Cloud Run, Firestore, Cloud SQL, BigQuery, Pub/Sub, GKE, Cloud Storage
Azure: Functions, App Service, Cosmos DB, SQL Database, Application Insights, AKS, Storage

Keep node labels SHORT (max 15 chars) - use abbreviations:
- "API GW" not "API Gateway"
- "Lambda" not "AWS Lambda Functions"
- "DDB" not "DynamoDB Database"
"""
                
                response = requests.post(
                    url="https://openrouter.ai/api/v1/chat/completions",
                    headers={
                        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
                        "Content-Type": "application/json",
                        "HTTP-Referer": "http://localhost:8501",
                        "X-Title": "ArchViz AI Pro",
                    },
                    data=json.dumps({
                        "model": "arcee-ai/trinity-large-preview:free",
                        "messages": [
                            {
                                "role": "user",
                                "content": f"""You are a cloud architecture diagram expert. Create a professional Mermaid diagram.

{icon_guidance}

REQUIREMENTS:
1. Use flowchart TB (top to bottom) or LR (left to right)
2. Keep ALL labels SHORT (under 15 characters)
3. Use exact service names from the list above
4. Style nodes with colors: compute (lightgreen), database (lightblue), storage (lightyellow), API (lightcoral), monitoring (lightgray)
5. Use arrows: --> for data flow, -.-> for async, ==> for main flow
6. Group services in subgraphs
7. Add clear flow labels

EXAMPLE:
graph TB
    subgraph "Client"
        A[Mobile App]
    end
    subgraph "API Layer"
        B[API Gateway]
        C[Lambda]
    end
    subgraph "Data"
        D[(DynamoDB)]
    end
    subgraph "Logs"
        E[CloudWatch]
    end

    A -->|"HTTPS"| B
    B -->|"Invoke"| C
    C -->|"Store"| D
    C -.->|"Logs"| E

    style B fill:#ffe1e1,stroke:#d32f2f,stroke-width:2px
    style C fill:#e1ffe1,stroke:#388e3c,stroke-width:2px
    style D fill:#e1e1ff,stroke:#303f9f,stroke-width:2px

Convert this to a diagram:
{prompt_text}

Output ONLY Mermaid code. No explanations."""
                            }
                        ]
                    })
                )
                
                result = response.json()
                
                if "choices" in result:
                    raw_mermaid = result['choices'][0]['message']['content'].strip()
                    clean_mermaid = raw_mermaid.replace("```mermaid", "").replace("```", "").strip()
                    
                    # Extract service names from the generated diagram
                    st.session_state['mermaid_code'] = clean_mermaid
                    st.session_state['show_editor'] = False
                    
                    st.success("✅ Diagram generated with official cloud service icons!")
                    st.rerun()
                else:
                    st.error(f"API Error: {result.get('error', 'Unknown error')}")
                    
            except Exception as e:
                st.error(f"Connection Error: {e}")

# Diagram Rendering
with col2:
    st.subheader("📊 Visual Diagram")

    # Icon legend
    st.markdown("""
    <div style="background: #f0f8ff; padding: 10px; border-radius: 5px; margin-bottom: 10px; font-size: 12px;">
        <b>🎨 Official Cloud Icons:</b>
        <span style="background: #e1ffe1; padding: 2px 8px; border-radius: 3px; margin: 0 5px;">AWS</span>
        <span style="background: #e3f2fd; padding: 2px 8px; border-radius: 3px; margin: 0 5px;">GCP</span>
        <span style="background: #e1f5fe; padding: 2px 8px; border-radius: 3px; margin: 0 5px;">Azure</span>
        | 
        <span style="background: #ffe1e1; padding: 2px 8px; border-radius: 3px; margin: 0 3px;">API</span>
        <span style="background: #e1ffe1; padding: 2px 8px; border-radius: 3px; margin: 0 3px;">Compute</span>
        <span style="background: #e1e1ff; padding: 2px 8px; border-radius: 3px; margin: 0 3px;">Database</span>
    </div>
    """, unsafe_allow_html=True)

    if 'mermaid_code' in st.session_state:
        mermaid_code = st.session_state['mermaid_code']

        # Build comprehensive icon mapping for JavaScript with PROVIDER-SPECIFIC matching
        icon_map_items = []

        # Create provider-specific service mappings with exact matches
        provider_services = {}
        
        # AWS services with their specific identifiers
        aws_services = {
            "lambda": ["lambda", "λ"],
            "api gateway": ["apigateway", "apigw", "apigateway", "api_gateway"],
            "dynamodb": ["dynamodb", "ddb"],
            "s3": ["s3"],
            "cloudwatch": ["cloudwatch", "cw"],
            "ec2": ["ec2"],
            "rds": ["rds"],
            "ecs": ["ecs"],
            "sqs": ["sqs"],
            "sns": ["sns"],
            "alb": ["alb", "applicationloadbalancer"],
            "elb": ["elb", "loadbalancer"],
            "cloudfront": ["cloudfront", "cf"],
            "vpc": ["vpc"],
            "route53": ["route53", "dns"],
            "cognito": ["cognito"],
            "kinesis": ["kinesis"],
            "step functions": ["stepfunctions", "stepfunction"],
            "eks": ["eks"],
            "fargate": ["fargate"],
            "elasticache": ["elasticache"],
            "aurora": ["aurora"]
        }
        
        # GCP services with their specific identifiers
        gcp_services = {
            "cloud functions": ["cloudfunctions", "gcf"],
            "cloud run": ["cloudrun"],
            "firestore": ["firestore"],
            "cloud sql": ["cloudsql"],
            "bigquery": ["bigquery", "bq"],
            "pub/sub": ["pubsub", "pubsub"],
            "gke": ["gke"],
            "cloud storage": ["cloudstorage", "gcs"],
            "cloud spanner": ["spanner"],
            "cloud tasks": ["cloudtasks"],
            "cloud scheduler": ["cloudscheduler"],
            "dataflow": ["dataflow"],
            "dataproc": ["dataproc"],
            "compute engine": ["computeengine", "gce"]
        }
        
        # Azure services with their specific identifiers
        azure_services = {
            "functions": ["azurefunctions"],
            "app service": ["appservice", "webapp"],
            "cosmos db": ["cosmosdb", "cosmos"],
            "sql database": ["sqldatabase", "azuresql"],
            "application insights": ["applicationinsights", "appinsights"],
            "aks": ["aks"],
            "storage": ["azurestorage", "blobstorage"],
            "logic apps": ["logicapps"],
            "service bus": ["servicebus"],
            "event hubs": ["eventhubs"],
            "key vault": ["keyvault"],
            "api management": ["apimanagement", "apim"],
            "virtual machines": ["virtualmachines", "azurevm"]
        }

        # Build the icon mapping with provider prefixes for disambiguation
        for service, url in SERVICE_ICONS.get("aws", {}).items():
            service_lower = service.lower()
            # Add with AWS prefix for precise matching
            icon_map_items.append(f'"aws_{service_lower.replace(" ", "").replace("-", "")}": "{url}"')
            
            # Add common variations
            if service_lower in aws_services:
                for variant in aws_services[service_lower]:
                    icon_map_items.append(f'"aws_{variant}": "{url}"')

        for service, url in SERVICE_ICONS.get("gcp", {}).items():
            service_lower = service.lower()
            # Add with GCP prefix for precise matching
            icon_map_items.append(f'"gcp_{service_lower.replace(" ", "").replace("-", "")}": "{url}"')
            
            # Add common variations
            if service_lower in gcp_services:
                for variant in gcp_services[service_lower]:
                    icon_map_items.append(f'"gcp_{variant}": "{url}"')

        for service, url in SERVICE_ICONS.get("azure", {}).items():
            service_lower = service.lower()
            # Add with Azure prefix for precise matching
            icon_map_items.append(f'"azure_{service_lower.replace(" ", "").replace("-", "")}": "{url}"')
            
            # Add common variations
            if service_lower in azure_services:
                for variant in azure_services[service_lower]:
                    icon_map_items.append(f'"azure_{variant}": "{url}"')

        # Generic icons (no provider prefix)
        for service, url in SERVICE_ICONS.get("generic", {}).items():
            key = service.lower().replace(" ", "").replace("-", "")
            icon_map_items.append(f'"{key}": "{url}"')

        icon_mapping_json = "{" + ",".join(icon_map_items) + "}"

        html_content = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <meta charset="UTF-8">
            <style>
                * {{
                    box-sizing: border-box;
                }}
                
                html, body {{
                    margin: 0;
                    padding: 0;
                    width: 100%;
                    height: 100%;
                }}
                
                body {{
                    padding: 20px;
                    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
                    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
                }}
                
                .diagram-container {{
                    background: white;
                    padding: 40px;
                    border-radius: 15px;
                    box-shadow: 0 8px 16px rgba(0,0,0,0.1);
                    max-width: 100%;
                    overflow: auto;
                    position: relative;
                }}
                .controls {{
                    position: absolute;
                    top: 10px;
                    right: 10px;
                    display: flex;
                    gap: 8px;
                    z-index: 100;
                    flex-wrap: wrap;
                }}
                .btn {{
                    background: #1976d2;
                    color: white;
                    border: none;
                    padding: 8px 12px;
                    border-radius: 5px;
                    cursor: pointer;
                    font-size: 13px;
                    display: flex;
                    align-items: center;
                    gap: 5px;
                    transition: all 0.2s;
                    white-space: nowrap;
                }}
                .btn:hover {{
                    background: #1565c0;
                    transform: translateY(-1px);
                }}
                .btn.download {{
                    background: #28a745;
                }}
                .btn.download:hover {{
                    background: #218838;
                }}
                .btn.fullscreen {{
                    background: #6c757d;
                }}
                .btn.fullscreen:hover {{
                    background: #5a6268;
                }}
                .zoom-level {{
                    background: #f0f0f0;
                    padding: 8px 12px;
                    border-radius: 5px;
                    font-size: 13px;
                    font-weight: bold;
                    min-width: 45px;
                    text-align: center;
                }}
                .mermaid-wrapper {{
                    transform-origin: center center;
                    transition: transform 0.3s ease;
                    display: flex;
                    justify-content: center;
                    align-items: center;
                    min-height: 400px;
                    padding: 30px;
                    position: relative;
                }}
                .mermaid {{
                    display: inline-block;
                    position: relative;
                }}

                /* Service Icon Styling */
                .icon-label-wrapper {{
                    display: flex !important;
                    flex-direction: column !important;
                    align-items: center !important;
                    gap: 4px !important;
                    padding: 4px !important;
                }}
                .icon-label-wrapper img {{
                    width: 32px !important;
                    height: 32px !important;
                    display: block !important;
                }}
                .icon-label-wrapper span {{
                    font-size: 11px !important;
                    text-align: center !important;
                    white-space: nowrap !important;
                    font-weight: 500 !important;
                }}

                foreignObject {{
                    overflow: visible !important;
                }}
                .nodeLabel {{
                    overflow: visible !important;
                }}

                /* Fullscreen styling */
                .diagram-container:fullscreen {{
                    background: white !important;
                    display: flex;
                    justify-content: center;
                    align-items: center;
                    padding: 80px 20px 100px 20px;
                }}
                
                .diagram-container:-webkit-full-screen {{
                    background: white !important;
                    display: flex;
                    justify-content: center;
                    align-items: center;
                    padding: 80px 20px 100px 20px;
                }}
                
                .diagram-container:-moz-full-screen {{
                    background: white !important;
                    display: flex;
                    justify-content: center;
                    align-items: center;
                    padding: 80px 20px 100px 20px;
                }}

                .mermaid-wrapper:fullscreen,
                .mermaid-wrapper:-webkit-full-screen,
                .mermaid-wrapper:-moz-full-screen {{
                    background: white;
                    border-radius: 15px;
                    padding: 50px;
                    max-width: 95vw;
                    max-height: 95vh;
                    overflow: auto;
                }}

                .fullscreen-controls {{
                    display: none;
                }}
                
                .diagram-container:fullscreen .fullscreen-controls,
                .diagram-container:-webkit-full-screen .fullscreen-controls,
                .diagram-container:-moz-full-screen .fullscreen-controls {{
                    display: flex;
                    position: fixed;
                    top: 20px;
                    right: 20px;
                    gap: 10px;
                    z-index: 10000;
                }}
                
                .fullscreen-zoom-controls {{
                    display: none;
                }}
                
                .diagram-container:fullscreen .fullscreen-zoom-controls,
                .diagram-container:-webkit-full-screen .fullscreen-zoom-controls,
                .diagram-container:-moz-full-screen .fullscreen-zoom-controls {{
                    display: flex;
                    position: fixed;
                    bottom: 20px;
                    left: 50%;
                    transform: translateX(-50%);
                    gap: 10px;
                    background: rgba(255,255,255,0.95);
                    padding: 10px 20px;
                    border-radius: 30px;
                    box-shadow: 0 4px 15px rgba(0,0,0,0.2);
                    z-index: 10000;
                }}

                .close-btn {{
                    background: #dc3545;
                    color: white;
                    border: none;
                    padding: 12px 20px;
                    border-radius: 8px;
                    cursor: pointer;
                    font-size: 16px;
                    font-weight: bold;
                    display: flex;
                    align-items: center;
                    gap: 8px;
                }}
                .close-btn:hover {{
                    background: #c82333;
                }}

                /* Enhanced Mermaid Styling */
                .node rect, .node circle, .node ellipse, .node polygon {{
                    stroke-width: 2px;
                    filter: drop-shadow(2px 2px 3px rgba(0,0,0,0.2));
                }}
                .edgePath path {{
                    stroke-width: 2px;
                }}
                .edgeLabel {{
                    background-color: white;
                    padding: 5px 10px;
                    border-radius: 5px;
                    font-size: 13px;
                    font-weight: 500;
                }}

                .edgeLabel, .edgeLabel span, .edgeLabel p {{
                    overflow: visible !important;
                    white-space: nowrap !important;
                    text-overflow: clip !important;
                    max-width: none !important;
                    width: auto !important;
                }}
                .edgeLabel foreignObject {{
                    overflow: visible !important;
                    width: auto !important;
                    min-width: 100px;
                }}
                .edgeLabel foreignObject div {{
                    overflow: visible !important;
                    white-space: nowrap !important;
                }}
                .edgeLabel rect {{
                    fill: white !important;
                    opacity: 0.9;
                }}

                .cluster rect {{
                    fill: #f9f9f9;
                    stroke: #666;
                    stroke-width: 2px;
                    stroke-dasharray: 5, 5;
                    rx: 10px;
                }}

                .nodeLabel {{
                    overflow: visible !important;
                    white-space: nowrap !important;
                }}
            </style>
        </head>
        <body>
            <div class="diagram-container" id="diagramContainer">
                <div class="controls">
                    <button class="btn" onclick="zoomOut()" title="Zoom Out">−</button>
                    <span class="zoom-level" id="zoomLevel">100%</span>
                    <button class="btn" onclick="zoomIn()" title="Zoom In">+</button>
                    <button class="btn" onclick="resetZoom()" title="Reset Zoom">↻</button>
                    <button class="btn download" onclick="downloadDiagram()" title="Download as PNG">⬇ Download PNG</button>
                    <button class="btn fullscreen" onclick="toggleFullscreen()" title="View Fullscreen (Press F)">⛶ Fullscreen</button>
                </div>
                
                <div class="fullscreen-controls">
                    <button class="btn download" onclick="downloadDiagram()">⬇ Download PNG</button>
                    <button class="close-btn" onclick="exitFullscreen()">✕ Close (Esc)</button>
                </div>
                
                <div class="mermaid-wrapper" id="mermaidWrapper">
                    <div class="mermaid" id="mainDiagram">
{mermaid_code}
                    </div>
                </div>
                
                <div class="fullscreen-zoom-controls">
                    <button class="btn" onclick="zoomOut()">−</button>
                    <span class="zoom-level" id="fullscreenZoomLevel">100%</span>
                    <button class="btn" onclick="zoomIn()">+</button>
                    <button class="btn" onclick="fitToWindow()">⛶ Fit</button>
                    <button class="btn" onclick="resetZoom()">↻</button>
                </div>
            </div>

            <script type="module">
                import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@11/dist/mermaid.esm.min.mjs';

                // Service icon mapping with provider-specific keys
                const serviceIcons = {icon_mapping_json};
                window.serviceIcons = serviceIcons;

                mermaid.initialize({{
                    startOnLoad: false,
                    theme: 'default',
                    flowchart: {{
                        useMaxWidth: false,
                        htmlLabels: true,
                        curve: 'basis',
                        padding: 30,
                        nodeSpacing: 70,
                        rankSpacing: 100,
                        diagramPadding: 20
                    }},
                    themeVariables: {{
                        fontSize: '15px',
                        fontFamily: 'Segoe UI, Arial, sans-serif',
                        primaryColor: '#e3f2fd',
                        primaryTextColor: '#1a1a1a',
                        primaryBorderColor: '#1976d2',
                        lineColor: '#424242',
                        secondaryColor: '#fff3e0',
                        tertiaryColor: '#f1f8e9'
                    }}
                }});

                // Detect which cloud provider(s) are in use from the diagram code
                function detectCloudProvider(diagramCode) {{
                    const code = diagramCode.toLowerCase();
                    const providers = [];
                    
                    // AWS indicators
                    const awsKeywords = ['lambda', 'dynamodb', 's3', 'ec2', 'rds', 'cloudwatch', 'api gateway', 
                                        'alb', 'elb', 'cloudfront', 'route53', 'cognito', 'kinesis', 'sns', 'sqs'];
                    if (awsKeywords.some(kw => code.includes(kw))) {{
                        providers.push('aws');
                    }}
                    
                    // GCP indicators
                    const gcpKeywords = ['cloud functions', 'cloud run', 'firestore', 'bigquery', 'pub/sub', 
                                        'gke', 'cloud sql', 'gcs', 'cloud storage', 'gce'];
                    if (gcpKeywords.some(kw => code.includes(kw))) {{
                        providers.push('gcp');
                    }}
                    
                    // Azure indicators  
                    const azureKeywords = ['app service', 'cosmos db', 'azure functions', 'aks', 'azure sql',
                                          'logic apps', 'service bus', 'event hubs', 'key vault'];
                    if (azureKeywords.some(kw => code.includes(kw))) {{
                        providers.push('azure');
                    }}
                    
                    return providers;
                }}

                // Enhanced icon finding with provider context
                function findIconForText(text, diagramCode) {{
                    const normalized = text.toLowerCase().replace(/[\\s\\-_\\/()\\[\\]]/g, '');
                    const providers = detectCloudProvider(diagramCode);
                    
                    console.log('Finding icon for:', text, 'Detected providers:', providers);
                    
                    // Priority 1: Exact matches with provider prefix
                    for (const provider of providers) {{
                        const providerKey = `${{provider}}_${{normalized}}`;
                        if (window.serviceIcons[providerKey]) {{
                            console.log('Found exact match:', providerKey);
                            return window.serviceIcons[providerKey];
                        }}
                    }}
                    
                    // Priority 2: Try common abbreviations with provider prefix
                    const commonMappings = {{
                        'apigw': ['apigateway', 'api_gateway'],
                        'alb': ['applicationloadbalancer'],
                        'ddb': ['dynamodb'],
                        'cw': ['cloudwatch'],
                        'cf': ['cloudfront'],
                        'bq': ['bigquery'],
                        'gcf': ['cloudfunctions'],
                        'gcs': ['cloudstorage']
                    }};
                    
                    for (const provider of providers) {{
                        if (commonMappings[normalized]) {{
                            for (const variant of commonMappings[normalized]) {{
                                const variantKey = `${{provider}}_${{variant}}`;
                                if (window.serviceIcons[variantKey]) {{
                                    console.log('Found variant match:', variantKey);
                                    return window.serviceIcons[variantKey];
                                }}
                            }}
                        }}
                    }}
                    
                    // Priority 3: Partial matching within same provider only
                    for (const provider of providers) {{
                        const providerPrefix = `${{provider}}_`;
                        for (const [key, url] of Object.entries(window.serviceIcons)) {{
                            if (key.startsWith(providerPrefix)) {{
                                const serviceKey = key.substring(providerPrefix.length);
                                // Only match if at least 5 characters match
                                if (normalized.length >= 5 && serviceKey.length >= 5) {{
                                    if (normalized.includes(serviceKey) || serviceKey.includes(normalized)) {{
                                        console.log('Found partial match:', key);
                                        return url;
                                    }}
                                }}
                            }}
                        }}
                    }}
                    
                    // Priority 4: Generic icons (no provider prefix)
                    if (window.serviceIcons[normalized]) {{
                        console.log('Found generic match:', normalized);
                        return window.serviceIcons[normalized];
                    }}
                    
                    console.log('No icon found for:', text);
                    return null;
                }}

                async function renderWithIcons() {{
                    const element = document.getElementById('mainDiagram');
                    const code = element.textContent;

                    try {{
                        const {{ svg }} = await mermaid.render('mermaid-svg', code);
                        element.innerHTML = svg;

                        setTimeout(() => {{
                            fixEdgeLabels('mainDiagram');
                            addIconsToNodes('mainDiagram', code);
                        }}, 200);
                        setTimeout(() => {{
                            fixEdgeLabels('mainDiagram');
                            addIconsToNodes('mainDiagram', code);
                        }}, 500);
                        setTimeout(() => {{
                            fixEdgeLabels('mainDiagram');
                            addIconsToNodes('mainDiagram', code);
                        }}, 1000);
                    }} catch (e) {{
                        console.error('Mermaid render error:', e);
                        element.innerHTML = '<p style="color: red;">Error rendering diagram. Please try regenerating.</p>';
                    }}
                }}

                function fixEdgeLabels(containerId) {{
                    const container = document.getElementById(containerId);
                    const svg = container.querySelector('svg');
                    if (!svg) return;

                    const edgeLabels = svg.querySelectorAll('.edgeLabel');

                    edgeLabels.forEach(label => {{
                        const fo = label.querySelector('foreignObject');
                        if (fo) {{
                            const textSpan = fo.querySelector('span, p, div');
                            if (textSpan) {{
                                const text = textSpan.textContent;
                                const tempSpan = document.createElement('span');
                                tempSpan.style.cssText = 'visibility: hidden; position: absolute; white-space: nowrap; font-size: 13px; font-family: inherit;';
                                tempSpan.textContent = text;
                                document.body.appendChild(tempSpan);
                                const textWidth = tempSpan.offsetWidth + 20;
                                document.body.removeChild(tempSpan);

                                fo.setAttribute('width', Math.max(textWidth, 50));
                                fo.setAttribute('height', '30');
                                fo.style.overflow = 'visible';

                                textSpan.style.whiteSpace = 'nowrap';
                                textSpan.style.overflow = 'visible';
                                textSpan.style.display = 'inline-block';
                                textSpan.style.width = 'auto';

                                const parentDiv = textSpan.closest('div');
                                if (parentDiv) {{
                                    parentDiv.style.overflow = 'visible';
                                    parentDiv.style.whiteSpace = 'nowrap';
                                    parentDiv.style.width = 'auto';
                                }}
                            }}
                        }}

                        const rect = label.querySelector('rect');
                        if (rect && fo) {{
                            const foWidth = fo.getAttribute('width');
                            if (foWidth) {{
                                rect.setAttribute('width', foWidth);
                            }}
                        }}
                    }});
                }}

                function addIconsToNodes(containerId, diagramCode) {{
                    const container = document.getElementById(containerId);
                    const svg = container.querySelector('svg');
                    if (!svg) return;

                    let iconsAdded = 0;
                    const foreignObjects = svg.querySelectorAll('foreignObject');

                    foreignObjects.forEach(fo => {{
                        const parentG = fo.closest('g');
                        if (parentG && (
                            parentG.classList.contains('cluster') ||
                            parentG.classList.contains('cluster-label') ||
                            parentG.querySelector('.cluster') ||
                            fo.closest('.cluster-label') ||
                            fo.closest('g.cluster')
                        )) {{
                            return;
                        }}

                        const nodeParent = fo.closest('g.node') || fo.closest('g.nodes') || fo.closest('[class*="node"]');
                        if (!nodeParent && parentG && !parentG.id?.includes('node')) {{
                            const nearbyCluster = svg.querySelector('.cluster');
                            if (nearbyCluster) {{
                                const isInsideNode = fo.closest('g.node') || fo.closest('g[class*="flowchart-label"]');
                                if (!isInsideNode) return;
                            }}
                        }}

                        const labelDiv = fo.querySelector('.nodeLabel');
                        if (!labelDiv) return;

                        const span = labelDiv.querySelector('span') || labelDiv;
                        const originalText = span.textContent.trim();

                        const nodeGroup = fo.closest('g.node') || fo.closest('g.nodes');
                        const hasNodeId = parentG && parentG.id && (parentG.id.includes('flowchart-') || parentG.id.includes('node'));
                        const hasShape = parentG && (
                            parentG.querySelector('rect:not(.cluster)') ||
                            parentG.querySelector('polygon') ||
                            parentG.querySelector('circle') ||
                            parentG.querySelector('ellipse') ||
                            parentG.querySelector('path.node')
                        );

                        if (!nodeGroup && !hasNodeId && !hasShape) return;

                        const commonSubgraphLabels = ['client', 'auth', 'api', 'data', 'monitoring', 'logs', 'storage', 'compute', 'network', 'security', 'frontend', 'backend', 'database', 'cache', 'web', 'mobile', 'users', 'services'];
                        if (commonSubgraphLabels.includes(originalText.toLowerCase()) && !hasShape) return;

                        if (labelDiv.querySelector('.service-icon-inline')) return;

                        const iconUrl = findIconForText(originalText, diagramCode);

                        if (iconUrl) {{
                            const wrapper = document.createElement('div');
                            wrapper.className = 'icon-label-wrapper';
                            wrapper.style.cssText = 'display: flex; flex-direction: column; align-items: center; gap: 2px; padding: 2px;';

                            const img = document.createElement('img');
                            img.src = iconUrl;
                            img.className = 'service-icon-inline';
                            img.style.cssText = 'width: 28px; height: 28px; display: block;';
                            img.onerror = function() {{ 
                                console.warn('Failed to load icon:', iconUrl);
                                this.style.display = 'none'; 
                            }};
                            img.onload = function() {{
                                console.log('Icon loaded successfully for:', originalText);
                                const foEl = this.closest('foreignObject');
                                if (foEl) {{
                                    const wrapperEl = this.closest('.icon-label-wrapper');
                                    if (wrapperEl) {{
                                        const rect = wrapperEl.getBoundingClientRect();
                                        foEl.setAttribute('width', Math.max(rect.width + 10, 60));
                                        foEl.setAttribute('height', Math.max(rect.height + 10, 55));
                                    }}
                                }}
                            }};

                            const textEl = document.createElement('span');
                            textEl.textContent = originalText;
                            textEl.style.cssText = 'font-size: 11px; text-align: center; white-space: nowrap; font-weight: 500;';

                            wrapper.appendChild(img);
                            wrapper.appendChild(textEl);

                            span.innerHTML = '';
                            span.appendChild(wrapper);

                            fo.setAttribute('width', Math.max(parseInt(fo.getAttribute('width')) || 50, 70));
                            fo.setAttribute('height', Math.max(parseInt(fo.getAttribute('height')) || 30, 55));

                            iconsAdded++;
                        }}
                    }});

                    console.log('Total icons added:', iconsAdded);
                }}

                window.addIconsToNodes = addIconsToNodes;
                renderWithIcons();
            </script>
            <script>
                let currentZoom = 1;
                const zoomStep = 0.25;
                const minZoom = 0.5;
                const maxZoom = 3;

                function updateZoom() {{
                    document.getElementById('mermaidWrapper').style.transform = `scale(${{currentZoom}})`;
                    document.getElementById('zoomLevel').textContent = Math.round(currentZoom * 100) + '%';
                    const fsZoomLevel = document.getElementById('fullscreenZoomLevel');
                    if (fsZoomLevel) {{
                        fsZoomLevel.textContent = Math.round(currentZoom * 100) + '%';
                    }}
                }}

                function zoomIn() {{
                    if (currentZoom < maxZoom) {{
                        currentZoom += zoomStep;
                        updateZoom();
                    }}
                }}

                function zoomOut() {{
                    if (currentZoom > minZoom) {{
                        currentZoom -= zoomStep;
                        updateZoom();
                    }}
                }}

                function resetZoom() {{
                    currentZoom = 1;
                    updateZoom();
                }}

                function fitToWindow() {{
                    const wrapper = document.getElementById('mermaidWrapper');
                    const svg = wrapper.querySelector('svg');
                    
                    if (svg && document.fullscreenElement) {{
                        wrapper.style.transform = 'scale(1)';
                        
                        setTimeout(() => {{
                            const svgRect = svg.getBoundingClientRect();
                            const containerWidth = window.innerWidth - 200;
                            const containerHeight = window.innerHeight - 200;

                            const scaleX = containerWidth / svgRect.width;
                            const scaleY = containerHeight / svgRect.height;
                            currentZoom = Math.min(scaleX, scaleY, 2);
                            currentZoom = Math.max(currentZoom, 0.3);
                            
                            updateZoom();
                        }}, 50);
                    }}
                }}

                function toggleFullscreen() {{
                    const container = document.getElementById('diagramContainer');
                    
                    if (!document.fullscreenElement) {{
                        if (container.requestFullscreen) {{
                            container.requestFullscreen();
                        }} else if (container.webkitRequestFullscreen) {{
                            container.webkitRequestFullscreen();
                        }} else if (container.mozRequestFullScreen) {{
                            container.mozRequestFullScreen();
                        }} else if (container.msRequestFullscreen) {{
                            container.msRequestFullscreen();
                        }}
                    }} else {{
                        exitFullscreen();
                    }}
                }}

                function exitFullscreen() {{
                    if (document.exitFullscreen) {{
                        document.exitFullscreen();
                    }} else if (document.webkitExitFullscreen) {{
                        document.webkitExitFullscreen();
                    }} else if (document.mozCancelFullScreen) {{
                        document.mozCancelFullScreen();
                    }} else if (document.msExitFullscreen) {{
                        document.msExitFullscreen();
                    }}
                }}

                document.addEventListener('fullscreenchange', function() {{
                    if (document.fullscreenElement) {{
                        setTimeout(fitToWindow, 300);
                    }} else {{
                        resetZoom();
                    }}
                }});

                async function urlToDataUrl(url) {{
                    try {{
                        const response = await fetch(url);
                        const blob = await response.blob();
                        return await new Promise(resolve => {{
                            const reader = new FileReader();
                            reader.onload = () => resolve(reader.result);
                            reader.readAsDataURL(blob);
                        }});
                    }} catch (e) {{
                        return null;
                    }}
                }}

                async function downloadDiagram() {{
                    const svg = document.querySelector('#mainDiagram svg');
                    if (!svg) {{
                        alert('Please generate a diagram first!');
                        return;
                    }}

                    const loadingDiv = document.createElement('div');
                    loadingDiv.style.cssText = 'position: fixed; top: 50%; left: 50%; transform: translate(-50%, -50%); background: rgba(0,0,0,0.8); color: white; padding: 20px 40px; border-radius: 10px; z-index: 99999;';
                    loadingDiv.textContent = 'Preparing download...';
                    document.body.appendChild(loadingDiv);

                    try {{
                        const clonedSvg = svg.cloneNode(true);
                        const bbox = svg.getBBox();

                        const padding = 100;
                        const width = bbox.width + (padding * 2);
                        const height = bbox.height + (padding * 2);

                        clonedSvg.setAttribute('width', width);
                        clonedSvg.setAttribute('height', height);
                        clonedSvg.setAttribute('viewBox', `${{bbox.x - padding}} ${{bbox.y - padding}} ${{width}} ${{height}}`);

                        const svgImages = clonedSvg.querySelectorAll('image');
                        for (const img of svgImages) {{
                            const href = img.getAttributeNS("http://www.w3.org/1999/xlink", "href") || img.getAttribute("href");
                            if (href && href.startsWith('http')) {{
                                const dataUrl = await urlToDataUrl(href);
                                if (dataUrl) {{
                                    img.setAttributeNS("http://www.w3.org/1999/xlink", "href", dataUrl);
                                }}
                            }}
                        }}

                        const htmlImages = clonedSvg.querySelectorAll('foreignObject img');
                        for (const img of htmlImages) {{
                            const src = img.getAttribute('src');
                            if (src && src.startsWith('http')) {{
                                const dataUrl = await urlToDataUrl(src);
                                if (dataUrl) {{
                                    img.setAttribute('src', dataUrl);
                                }}
                            }}
                        }}

                        const svgData = new XMLSerializer().serializeToString(clonedSvg);

                        const scale = 3;
                        const canvas = document.createElement('canvas');
                        const ctx = canvas.getContext('2d');
                        canvas.width = width * scale;
                        canvas.height = height * scale;

                        ctx.fillStyle = 'white';
                        ctx.fillRect(0, 0, canvas.width, canvas.height);

                        const imgEl = new Image();
                        imgEl.onload = function() {{
                            ctx.drawImage(imgEl, 0, 0, canvas.width, canvas.height);
                            canvas.toBlob(function(blob) {{
                                const url = URL.createObjectURL(blob);
                                const a = document.createElement('a');
                                a.href = url;
                                a.download = 'architecture-diagram.png';
                                document.body.appendChild(a);
                                a.click();
                                document.body.removeChild(a);
                                URL.revokeObjectURL(url);
                                loadingDiv.remove();
                            }}, 'image/png');
                        }};
                        imgEl.onerror = function() {{
                            loadingDiv.remove();
                            alert('Failed to export diagram. Please try again.');
                        }};

                        imgEl.src = 'data:image/svg+xml;base64,' + btoa(unescape(encodeURIComponent(svgData)));
                    }} catch (e) {{
                        console.error('Download error:', e);
                        loadingDiv.remove();
                        alert('Error preparing download: ' + e.message);
                    }}
                }}

                document.addEventListener('keydown', function(e) {{
                    if (e.key === 'Escape' && document.fullscreenElement) {{
                        exitFullscreen();
                    }}
                    if ((e.key === 'f' || e.key === 'F') && !e.ctrlKey && !e.metaKey) {{
                        toggleFullscreen();
                        e.preventDefault();
                    }}
                }});

                window.zoomIn = zoomIn;
                window.zoomOut = zoomOut;
                window.resetZoom = resetZoom;
                window.downloadDiagram = downloadDiagram;
                window.toggleFullscreen = toggleFullscreen;
                window.exitFullscreen = exitFullscreen;
                window.fitToWindow = fitToWindow;
            </script>
        </body>
        </html>
        """
        components.html(html_content, height=700, scrolling=True)
        
        with st.expander("📄 View/Download Source Code"):
            st.code(mermaid_code, language="text")
            col_down1, col_down2, col_down3 = st.columns(3)
            with col_down1:
                st.download_button(
                    label="⬇️ Mermaid (.mmd)",
                    data=mermaid_code,
                    file_name="architecture-diagram.mmd",
                    mime="text/plain",
                    use_container_width=True
                )
            with col_down2:
                st.download_button(
                    label="⬇️ Text (.txt)",
                    data=mermaid_code,
                    file_name="architecture-diagram.txt",
                    mime="text/plain",
                    use_container_width=True
                )
            with col_down3:
                markdown_content = f"# Architecture Diagram\n\n```mermaid\n{mermaid_code}\n```"
                st.download_button(
                    label="⬇️ Markdown (.md)",
                    data=markdown_content,
                    file_name="architecture-diagram.md",
                    mime="text/markdown",
                    use_container_width=True
                )
            
            st.caption("💡 Import into: Mermaid Live • GitHub • Confluence • Notion • Draw.io")
    else:
        st.info("👈 Describe your architecture and click 'Generate Diagram' to see it with official cloud service icons!")
        
        st.markdown("### 🎨 Available Service Icons")
        
        tab_aws, tab_gcp, tab_azure = st.tabs(["AWS", "GCP", "Azure"])
        
        with tab_aws:
            st.markdown("**AWS Architecture Icons (Sample)**")
            sample_aws = ["Lambda", "API Gateway", "DynamoDB", "S3", "CloudWatch", "EC2", "RDS", "ECS"]
            icons_html = ""
            for service in sample_aws:
                icon_url = get_icon_url(service.lower(), "aws")
                if icon_url:
                    icons_html += f'<div style="display: inline-block; text-align: center; margin: 10px;"><img src="{icon_url}" width="48" height="48"><br><span style="font-size: 11px;">{service}</span></div>'
            st.markdown(icons_html, unsafe_allow_html=True)
        
        with tab_gcp:
            st.markdown("**Google Cloud Icons (Sample)**")
            sample_gcp = ["Cloud Functions", "Cloud Run", "Firestore", "Cloud SQL", "BigQuery", "Pub/Sub", "GKE", "Cloud Storage"]
            icons_html = ""
            for service in sample_gcp:
                icon_url = get_icon_url(service.lower(), "gcp")
                if icon_url:
                    icons_html += f'<div style="display: inline-block; text-align: center; margin: 10px;"><img src="{icon_url}" width="48" height="48"><br><span style="font-size: 11px;">{service}</span></div>'
            st.markdown(icons_html, unsafe_allow_html=True)
        
        with tab_azure:
            st.markdown("**Azure Service Icons (Sample)**")
            sample_azure = ["Functions", "App Service", "Cosmos DB", "SQL Database", "Application Insights", "AKS", "Storage", "Logic Apps"]
            icons_html = ""
            for service in sample_azure:
                icon_url = get_icon_url(service.lower(), "azure")
                if icon_url:
                    icons_html += f'<div style="display: inline-block; text-align: center; margin: 10px;"><img src="{icon_url}" width="48" height="48"><br><span style="font-size: 11px;">{service}</span></div>'
            st.markdown(icons_html, unsafe_allow_html=True)

st.markdown("---")
st.caption("© 2026 ArchViz AI Pro")