# 🚀 Modal Deployment Script (Marimo Configuration)
# Generated by modal-for-noobs - https://github.com/arthrod/modal-for-noobs
# Deployment Mode: marimo
# Features: Gradio app with Marimo notebooks, ML libraries, dashboard, and logging
# Timeout: {{ timeout_seconds }}s

import modal
import sys
import os
import asyncio
import subprocess
from pathlib import Path
from datetime import datetime
from fastapi import FastAPI, Request
from fastapi.responses import RedirectResponse, StreamingResponse
import gradio as gr
from gradio.routes import mount_gradio_app
from loguru import logger
import httpx

# Configuration constants
APP_NAME = "{{ app_name }}"
APP_TITLE = "{{ app_name }} - Modal Dashboard (Marimo)"
APP_DESCRIPTION = "Optimized deployment with Marimo notebooks and monitoring"
DEPLOYMENT_MODE = "marimo"
TIMEOUT_SECONDS = {{ timeout_seconds }}
MAX_CONTAINERS = {{ max_containers }}
MIN_CONTAINERS = {{ min_containers }}

# Configure logging
logger.add(sys.stderr, format="{time} {level} {message}", level="INFO")

# Create Modal App
app = modal.App(APP_NAME)

{% if provision_nfs %}
# Persistent storage for notebooks
WORKSPACE_PATH = "/workspace"
volume = modal.Volume.persisted(f"{APP_NAME}-marimo-workspace")
{% endif %}

# Container Image Configuration (Marimo + ML)
{{ image_config }}

# Original Application Code
{{ original_code }}

# Marimo notebook template
MARIMO_NOTEBOOK_TEMPLATE = '''import marimo

__generated_with = "0.1.0"
app = marimo.App()

@app.cell
def __():
    import marimo as mo
    return mo,

@app.cell
def __(mo):
    mo.md(r"""
    # Welcome to Modal Marimo! 🎉
    
    This Marimo notebook is running alongside your Gradio app on Modal.
    
    Marimo provides:
    - **Reactive notebooks** - cells automatically re-run when dependencies change
    - **Pure Python** - notebooks are just Python files
    - **Interactive widgets** - build UIs with Python
    - **GPU acceleration** - full ML/AI support
    
    Your Gradio app is available at the [root URL](/).
    """)
    return

@app.cell
def __():
    import sys
    import torch
    import pandas as pd
    import numpy as np
    
    print(f"Python {sys.version}")
    print(f"PyTorch {torch.__version__}")
    print(f"CUDA available: {torch.cuda.is_available()}")
    
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    return np, pd, sys, torch

@app.cell
def __(mo):
    mo.md("## Interactive Example")
    slider = mo.ui.slider(1, 10, value=5, label="Select a value")
    return slider,

@app.cell
def __(mo, slider):
    mo.md(f"You selected: **{slider.value}**")
    
    # Create a simple plot
    import matplotlib.pyplot as plt
    
    fig, ax = plt.subplots()
    ax.bar(range(slider.value), range(slider.value))
    ax.set_title(f"Bar chart with {slider.value} bars")
    plt.tight_layout()
    return ax, fig, plt

if __name__ == "__main__":
    app.run()
'''

async def start_marimo_server():
    """Start Marimo server in the background."""
    {% if provision_nfs %}
    # Use persistent workspace
    workspace_dir = Path(WORKSPACE_PATH)
    {% else %}
    # Create temporary workspace
    workspace_dir = Path("/tmp/marimo-workspace")
    {% endif %}
    workspace_dir.mkdir(exist_ok=True, parents=True)
    
    # Create a sample Marimo notebook
    sample_notebook = workspace_dir / "welcome.py"
    if not sample_notebook.exists():
        sample_notebook.write_text(MARIMO_NOTEBOOK_TEMPLATE)
    
    # Start Marimo server
    logger.info("Starting Marimo server on port 2718")
    process = await asyncio.create_subprocess_exec(
        "marimo", "run",
        "--host", "0.0.0.0",
        "--port", "2718",
        "--no-browser",
        str(workspace_dir),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    
    # Give it time to start
    await asyncio.sleep(3)
    logger.info("Marimo server started")
    
    return process

# Modal Function Configuration with GPU
@app.function(
    image=image,
    gpu="any",  # GPU support for ML workloads
    min_containers={{ min_containers }},
    max_containers={{ max_containers }},
    timeout={{ timeout_seconds }},
    memory=16384,  # 16GB RAM
    {% if provision_nfs %}
    volumes={WORKSPACE_PATH: volume},
    {% endif %}
)
@modal.concurrent(max_inputs=100)
@modal.asgi_app()
def deploy_gradio():
    """Deploy Gradio app with Marimo notebooks and dashboard on Modal."""
    
    # Check GPU availability
    import torch
    gpu_available = torch.cuda.is_available()
    gpu_name = torch.cuda.get_device_name(0) if gpu_available else "None"
    
    logger.info(f"Starting Modal deployment in marimo mode (GPU: {gpu_available})")
    
    # Start Marimo in the background
    asyncio.create_task(start_marimo_server())
    
    # Detect Gradio Interface
    demo = None
    interface_names = ['demo', 'app', 'interface', 'iface']
    
    for name in interface_names:
        if name in globals() and hasattr(globals()[name], 'launch'):
            demo = globals()[name]
            logger.info(f"Found Gradio interface: {name}")
            break
    
    if demo is None:
        for var_name, var_value in globals().items():
            if hasattr(var_value, 'queue') and hasattr(var_value, 'launch'):
                demo = var_value
                logger.info(f"Found Gradio interface through scanning: {var_name}")
                break
    
    if demo is None:
        logger.error("No Gradio interface found")
        raise ValueError("Could not find Gradio interface")
    
    # Create Dashboard with Marimo integration
    with gr.Blocks() as enhanced_dashboard:
        gr.Markdown("# 🚀 Modal Deployment Dashboard")
        gr.Markdown("Monitor and manage your Modal deployment with Marimo notebook integration")
        
        with gr.Tabs():
            # Marimo Tab (first for easy access)
            with gr.Tab("📓 Marimo Notebooks"):
                gr.Markdown("### Interactive Marimo Notebook Environment")
                gr.Markdown("Access Marimo for reactive Python notebooks with GPU support.")
                
                with gr.Row():
                    gr.Markdown("**Marimo is running on port 2718**")
                    marimo_btn = gr.Button("🚀 Open Marimo", variant="primary", scale=2)
                
                gr.Markdown("""
                #### Why Marimo?
                - **Reactive notebooks** - Cells automatically update when dependencies change
                - **Pure Python files** - Version control friendly, no JSON
                - **Built-in UI components** - Create interactive apps with Python
                - **GPU acceleration** - Full PyTorch and ML support
                
                #### Features:
                - Pre-installed ML/AI packages
                - GPU support enabled
                - Interactive widgets and visualizations
                {% if provision_nfs %}
                - Persistent workspace at `{{ WORKSPACE_PATH }}`
                {% else %}
                - Temporary workspace for this session
                {% endif %}
                
                #### Usage:
                1. Click the button above to open Marimo in a new tab
                2. Create reactive notebooks with `.py` extension
                3. Build interactive ML experiments
                """)
                
                # JavaScript to open Marimo in new tab
                marimo_btn.click(
                    None,
                    None,
                    None,
                    js="window.open('/marimo', '_blank')"
                )
            
            # App Tab
            with gr.Tab("🎯 Your App"):
                demo.render()
            
            # Status Tab
            with gr.Tab("📊 Status"):
                gr.Markdown("### Deployment Information")
                status_info = gr.Textbox(
                    value=f"""
                    **App Name:** {APP_NAME}
                    **Mode:** {DEPLOYMENT_MODE}
                    **GPU Enabled:** {gpu_available}
                    **GPU Name:** {gpu_name}
                    **Timeout:** {TIMEOUT_SECONDS}s
                    **Max Containers:** {MAX_CONTAINERS}
                    """,
                    label="Deployment Details",
                    lines=8,
                    interactive=False
                )
        
        gr.Markdown("---")
        gr.Markdown("🚀 Powered by [Modal](https://modal.com) | 📓 [Marimo](/marimo) | Generated by [modal-for-noobs](https://github.com/arthrod/modal-for-noobs)")
    
    enhanced_dashboard.queue(max_size=20)
    
    # FastAPI Setup with Marimo proxy
    fastapi_app = FastAPI(
        title=APP_TITLE,
        description=APP_DESCRIPTION,
        version="1.0.0",
        docs_url="/docs",
        redoc_url="/redoc"
    )
    
    # Marimo proxy client
    marimo_client = httpx.AsyncClient(base_url="http://localhost:2718")
    
    # Add Marimo proxy endpoints
    @fastapi_app.get("/marimo")
    @fastapi_app.get("/marimo/{path:path}")
    async def marimo_proxy(request: Request, path: str = ""):
        """Proxy requests to Marimo."""
        # Get the full URL with query params
        url = f"/{path}"
        if request.url.query:
            url += f"?{request.url.query}"
        
        # Proxy the request
        try:
            response = await marimo_client.request(
                method=request.method,
                url=url,
                headers={k: v for k, v in request.headers.items() if k.lower() not in ['host', 'connection']},
                content=await request.body() if request.method in ["POST", "PUT", "PATCH"] else None,
            )
            
            # Stream the response
            return StreamingResponse(
                response.aiter_bytes(),
                status_code=response.status_code,
                headers=dict(response.headers)
            )
        except Exception as e:
            logger.error(f"Marimo proxy error: {e}")
            return RedirectResponse(url="/")
    
    # Health check endpoint
    @fastapi_app.get("/health")
    async def health_check():
        return {
            "status": "healthy",
            "app_name": APP_NAME,
            "deployment_mode": DEPLOYMENT_MODE,
            "timestamp": datetime.now().isoformat(),
            "gpu_available": gpu_available,
            "marimo_enabled": True,
        }
    
    logger.info("Dashboard with Marimo integration configured successfully")
    
    # Mount Gradio app
    return mount_gradio_app(fastapi_app, enhanced_dashboard, path="/")

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