#!/usr/bin/env nu

# Marimo Interactive Dashboard Integration
# Creates interactive notebooks and dashboards for infrastructure monitoring

use ../dataframes/polars_integration.nu *
use ../observability/collectors.nu *
use ../observability/agents.nu *
use ../api/server.nu *

# Check if Marimo is available
export def check_marimo_available []: nothing -> bool {
    (which marimo | length > 0)
}

# Install Marimo if not available
export def install_marimo []: nothing -> bool {
    if not (check_marimo_available) {
        print "📦 Installing Marimo..."
        try {
            ^pip install marimo
            true
        } catch {
            print "❌ Failed to install Marimo. Please install manually: pip install marimo"
            false
        }
    } else {
        true
    }
}

# Create interactive dashboard
export def create_dashboard [
    --name: string = "infrastructure-dashboard"
    --data_sources: list<string> = ["logs", "metrics", "infrastructure"]
    --refresh_interval: duration = 30sec
    --port: int = 8080
]: nothing -> nothing {

    if not (install_marimo) {
        error make { msg: "Marimo installation failed" }
    }

    print $"🚀 Creating interactive dashboard: ($name)"

    # Generate dashboard Python file
    let dashboard_code = generate_dashboard_code $data_sources $refresh_interval
    let dashboard_path = $"dashboards/($name).py"

    # Create dashboards directory
    mkdir dashboards

    # Write dashboard file
    $dashboard_code | save --force $dashboard_path

    print $"📊 Dashboard created at: ($dashboard_path)"
    print $"🌐 Starting dashboard on port ($port)..."

    # Start Marimo dashboard
    ^marimo run $dashboard_path --port $port --host "0.0.0.0"
}

# Generate dashboard Python code
def generate_dashboard_code [
    data_sources: list<string>
    refresh_interval: duration
]: [list<string>, duration] -> string {

    let refresh_ms = ($refresh_interval | into int) / 1000000

    $"
import marimo as mo
import polars as pl
import plotly.graph_objects as go
import plotly.express as px
from datetime import datetime, timedelta
import asyncio
import requests
import json

# Configure the app
app = mo.App(width=\"full\")

@app.cell
def header():
    mo.md(
        '''
        # 🚀 Systems Provisioning Dashboard

        Real-time monitoring and analytics for your infrastructure
        '''
    )
    return

@app.cell
def data_sources_config():
    # Data source configuration
    DATA_SOURCES = ($data_sources | to json)
    REFRESH_INTERVAL = ($refresh_ms)
    API_BASE = \"http://localhost:3000\"
    return DATA_SOURCES, REFRESH_INTERVAL, API_BASE

@app.cell
def fetch_data(DATA_SOURCES, API_BASE):
    '''Fetch data from provisioning API'''

    def get_api_data(endpoint):
        try:
            response = requests.get(f\"{API_BASE}/api/{endpoint}\")
            return response.json() if response.status_code == 200 else {}
        except:
            return {}

    # Fetch data from different sources
    logs_data = get_api_data(\"logs\") if \"logs\" in DATA_SOURCES else {}
    metrics_data = get_api_data(\"metrics\") if \"metrics\" in DATA_SOURCES else {}
    infra_data = get_api_data(\"query/infrastructure\") if \"infrastructure\" in DATA_SOURCES else {}

    return logs_data, metrics_data, infra_data

@app.cell
def logs_analysis(logs_data):
    '''Analyze logs data'''
    if not logs_data:
        return mo.md(\"📝 No logs data available\")

    # Convert to DataFrame
    try:
        df_logs = pl.DataFrame(logs_data.get('logs', []))

        if df_logs.height == 0:
            return mo.md(\"📝 No log entries found\")

        # Log level distribution
        level_counts = df_logs.group_by(\"level\").agg(pl.count().alias(\"count\"))

        fig_levels = px.pie(
            level_counts.to_pandas(),
            values='count',
            names='level',
            title=\"Log Levels Distribution\"
        )

        # Recent errors
        if \"timestamp\" in df_logs.columns:
            recent_errors = df_logs.filter(
                pl.col(\"level\").is_in([\"error\", \"fatal\", \"warn\"])
            ).sort(\"timestamp\", descending=True).head(10)

            error_table = mo.ui.table(
                recent_errors.to_pandas(),
                selection=\"single\"
            )
        else:
            error_table = mo.md(\"No timestamp data available\")

        return mo.vstack([
            mo.md(\"## 📊 Logs Analysis\"),
            mo.ui.plotly(fig_levels),
            mo.md(\"### Recent Errors/Warnings\"),
            error_table
        ])

    except Exception as e:
        return mo.md(f\"❌ Error processing logs: {e}\")

@app.cell
def metrics_dashboard(metrics_data):
    '''System metrics dashboard'''
    if not metrics_data:
        return mo.md(\"📈 No metrics data available\")

    try:
        # System metrics visualization
        metrics = metrics_data.get('metrics', {})

        # CPU Usage
        cpu_data = metrics.get('cpu', {})
        if cpu_data:
            fig_cpu = go.Figure()
            fig_cpu.add_trace(go.Scatter(
                x=list(range(len(cpu_data.get('values', [])))),
                y=cpu_data.get('values', []),
                mode='lines+markers',
                name='CPU %',
                line=dict(color='#ff6b6b')
            ))
            fig_cpu.update_layout(title='CPU Usage Over Time', yaxis_title='Percentage')
        else:
            fig_cpu = None

        # Memory Usage
        memory_data = metrics.get('memory', {})
        if memory_data:
            fig_memory = go.Figure()
            fig_memory.add_trace(go.Scatter(
                x=list(range(len(memory_data.get('values', [])))),
                y=memory_data.get('values', []),
                mode='lines+markers',
                name='Memory %',
                line=dict(color='#4ecdc4')
            ))
            fig_memory.update_layout(title='Memory Usage Over Time', yaxis_title='Percentage')
        else:
            fig_memory = None

        # Infrastructure status
        infra_status = metrics.get('infrastructure', {})
        status_cards = []

        if infra_status:
            for service, data in infra_status.items():
                status = \"🟢 Healthy\" if data.get('healthy', False) else \"🔴 Unhealthy\"
                status_cards.append(
                    mo.md(f\"**{service}**: {status} (Load: {data.get('load', 'N/A')})\")
                )

        components = [mo.md(\"## 📈 System Metrics\")]

        if fig_cpu:
            components.append(mo.ui.plotly(fig_cpu))
        if fig_memory:
            components.append(mo.ui.plotly(fig_memory))

        if status_cards:
            components.extend([mo.md(\"### Infrastructure Status\")] + status_cards)

        return mo.vstack(components)

    except Exception as e:
        return mo.md(f\"❌ Error processing metrics: {e}\")

@app.cell
def infrastructure_overview(infra_data):
    '''Infrastructure overview and topology'''
    if not infra_data:
        return mo.md(\"🏗️ No infrastructure data available\")

    try:
        infra = infra_data.get('infrastructure', {})

        # Servers overview
        servers = infra.get('servers', [])
        if servers:
            df_servers = pl.DataFrame(servers)

            # Provider distribution
            if \"provider\" in df_servers.columns:
                provider_counts = df_servers.group_by(\"provider\").agg(pl.count().alias(\"count\"))
                fig_providers = px.bar(
                    provider_counts.to_pandas(),
                    x='provider',
                    y='count',
                    title='Servers by Provider'
                )
            else:
                fig_providers = None

            # Status distribution
            if \"status\" in df_servers.columns:
                status_counts = df_servers.group_by(\"status\").agg(pl.count().alias(\"count\"))
                fig_status = px.pie(
                    status_counts.to_pandas(),
                    values='count',
                    names='status',
                    title='Server Status Distribution'
                )
            else:
                fig_status = None

            # Server table
            server_table = mo.ui.table(
                df_servers.to_pandas(),
                selection=\"multiple\"
            )

            components = [
                mo.md(\"## 🏗️ Infrastructure Overview\"),
                mo.md(f\"**Total Servers**: {len(servers)}\")
            ]

            if fig_providers:
                components.append(mo.ui.plotly(fig_providers))
            if fig_status:
                components.append(mo.ui.plotly(fig_status))

            components.extend([
                mo.md(\"### Server Details\"),
                server_table
            ])

            return mo.vstack(components)
        else:
            return mo.md(\"🏗️ No server data available\")

    except Exception as e:
        return mo.md(f\"❌ Error processing infrastructure data: {e}\")

@app.cell
def ai_insights():
    '''AI-powered insights and recommendations'''

    # This would integrate with our AI agents
    insights = [
        \"💡 **Cost Optimization**: Consider downsizing instance i-12345 (38% CPU avg)\",
        \"⚠️ **Performance Alert**: Database response time increased 15% in last hour\",
        \"🔮 **Prediction**: Disk space on /var/log will be full in 3 days\",
        \"🛡️ **Security**: No failed login attempts detected in last 24h\",
        \"📈 **Scaling**: Web tier may need +2 instances based on traffic trends\"
    ]

    insight_cards = [mo.md(insight) for insight in insights]

    return mo.vstack([
        mo.md(\"## 🤖 AI Insights & Recommendations\"),
        mo.md(\"_Powered by Rust-based AI agents_\"),
        *insight_cards
    ])

@app.cell
def controls():
    '''Dashboard controls and settings'''

    refresh_button = mo.ui.button(
        label=\"🔄 Refresh Data\",
        on_click=lambda: print(\"Refreshing dashboard data...\")
    )

    auto_refresh = mo.ui.checkbox(
        label=\"Auto-refresh every 30 seconds\",
        value=True
    )

    export_button = mo.ui.button(
        label=\"📊 Export Report\",
        on_click=lambda: print(\"Exporting dashboard report...\")
    )

    return mo.hstack([refresh_button, auto_refresh, export_button])

@app.cell
def footer():
    mo.md(
        '''
        ---
        **Systems Provisioning Dashboard** | Powered by Rust + Nushell + Marimo
        🔗 [API Status](http://localhost:3000/health) | 📖 [Documentation](http://localhost:3000/docs)
        '''
    )
    return

if __name__ == \"__main__\":
    app.run()
"
}

# Create predefined dashboard templates
export def create_template [
    template: string
    --name: string = ""
]: string -> nothing {

    let dashboard_name = if ($name | is-empty) { $"($template)-dashboard" } else { $name }

    match $template {
        "monitoring" => {
            create_dashboard --name $dashboard_name --data_sources ["logs", "metrics"] --refresh_interval 15sec
        }
        "infrastructure" => {
            create_dashboard --name $dashboard_name --data_sources ["infrastructure", "metrics"] --refresh_interval 30sec
        }
        "full" => {
            create_dashboard --name $dashboard_name --data_sources ["logs", "metrics", "infrastructure"] --refresh_interval 30sec
        }
        "ai-insights" => {
            create_dashboard --name $dashboard_name --data_sources ["logs", "metrics", "infrastructure"] --refresh_interval 10sec
        }
        _ => {
            error make { msg: $"Unknown template: ($template). Available: monitoring, infrastructure, full, ai-insights" }
        }
    }
}

# List available dashboards
export def list_dashboards []: nothing -> list<record> {
    if not ("dashboards" | path exists) {
        return []
    }

    ls dashboards/*.py
    | get name
    | each {|path|
        {
            name: ($path | path basename | str replace ".py" "")
            path: $path
            size: (stat $path | get size)
            modified: (stat $path | get modified)
        }
    }
}

# Start existing dashboard
export def start_dashboard [
    dashboard_name: string
    --port: int = 8080
    --host: string = "0.0.0.0"
]: string -> nothing {

    let dashboard_path = $"dashboards/($dashboard_name).py"

    if not ($dashboard_path | path exists) {
        error make { msg: $"Dashboard not found: ($dashboard_path)" }
    }

    print $"🌐 Starting dashboard: ($dashboard_name) on ($host):($port)"
    ^marimo run $dashboard_path --port $port --host $host
}

# Export dashboard as static HTML
export def export_dashboard [
    dashboard_name: string
    --output: string = ""
]: string -> nothing {

    let dashboard_path = $"dashboards/($dashboard_name).py"
    let output_path = if ($output | is-empty) { $"exports/($dashboard_name).html" } else { $output }

    if not ($dashboard_path | path exists) {
        error make { msg: $"Dashboard not found: ($dashboard_path)" }
    }

    # Create exports directory
    mkdir exports

    print $"📤 Exporting dashboard to: ($output_path)"
    ^marimo export html $dashboard_path --output $output_path

    print $"✅ Dashboard exported successfully"
}

# Dashboard management commands
export def main [
    command: string
    ...args: string
]: [string, ...string] -> nothing {

    match $command {
        "create" => {
            if ($args | length) >= 1 {
                let template = $args.0
                let name = if ($args | length) >= 2 { $args.1 } else { "" }
                create_template $template --name $name
            } else {
                create_dashboard
            }
        }
        "list" => {
            list_dashboards | table
        }
        "start" => {
            if ($args | length) >= 1 {
                let name = $args.0
                let port = if ($args | length) >= 2 { $args.1 | into int } else { 8080 }
                start_dashboard $name --port $port
            } else {
                error make { msg: "Dashboard name required" }
            }
        }
        "export" => {
            if ($args | length) >= 1 {
                let name = $args.0
                let output = if ($args | length) >= 2 { $args.1 } else { "" }
                export_dashboard $name --output $output
            } else {
                error make { msg: "Dashboard name required" }
            }
        }
        "install" => {
            install_marimo
        }
        _ => {
            print "📊 Marimo Dashboard Integration Commands:"
            print ""
            print "Usage: marimo_integration <command> [args...]"
            print ""
            print "Commands:"
            print "  create [template] [name]  - Create new dashboard from template"
            print "  list                      - List available dashboards"
            print "  start <name> [port]       - Start existing dashboard"
            print "  export <name> [output]    - Export dashboard to HTML"
            print "  install                   - Install Marimo package"
            print ""
            print "Templates:"
            print "  monitoring    - Logs and metrics dashboard"
            print "  infrastructure- Infrastructure overview"
            print "  full          - Complete monitoring dashboard"
            print "  ai-insights   - AI-powered insights dashboard"
        }
    }
}