# 🚀 Modal Deployment Script (Gradio-Jupyter Configuration)
# Generated by modal-for-noobs - https://github.com/arthrod/modal-for-noobs
# Deployment Mode: gradio-jupyter
# Features: Gradio app with Jupyter notebook integration, ML libraries, and monitoring
# 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 (Jupyter)"
APP_DESCRIPTION = "Deployment with Jupyter notebooks and monitoring"
DEPLOYMENT_MODE = "gradio-jupyter"
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}-jupyter-workspace")
{% endif %}

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

# Original Application Code
{{ original_code }}

# Jupyter notebook template
JUPYTER_NOTEBOOK_TEMPLATE = {
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Welcome to Modal Jupyter! 🎉\\n",
                "\\n",
                "This Jupyter notebook is running alongside your Gradio app on Modal.\\n",
                "\\n",
                "## Features:\\n",
                "- **GPU acceleration** - Full ML/AI support\\n",
                "- **Pre-installed ML packages** - PyTorch, TensorFlow, scikit-learn, etc.\\n",
                "- **Persistent storage** - Your notebooks are saved\\n",
                "- **Integration with Gradio** - Access your app at the root URL\\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import sys\\n",
                "import torch\\n",
                "import pandas as pd\\n",
                "import numpy as np\\n",
                "\\n",
                "print(f'Python {sys.version}')\\n",
                "print(f'PyTorch {torch.__version__}')\\n",
                "print(f'CUDA available: {torch.cuda.is_available()}')\\n",
                "\\n",
                "if torch.cuda.is_available():\\n",
                "    print(f'GPU: {torch.cuda.get_device_name(0)}')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Example: Simple ML Model\\n",
                "\\n",
                "Let's create a simple neural network as an example:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Simple neural network example\\n",
                "import torch.nn as nn\\n",
                "import torch.optim as optim\\n",
                "\\n",
                "class SimpleNet(nn.Module):\\n",
                "    def __init__(self):\\n",
                "        super(SimpleNet, self).__init__()\\n",
                "        self.fc1 = nn.Linear(10, 50)\\n",
                "        self.fc2 = nn.Linear(50, 1)\\n",
                "        \\n",
                "    def forward(self, x):\\n",
                "        x = torch.relu(self.fc1(x))\\n",
                "        x = self.fc2(x)\\n",
                "        return x\\n",
                "\\n",
                "# Create model\\n",
                "model = SimpleNet()\\n",
                "if torch.cuda.is_available():\\n",
                "    model = model.cuda()\\n",
                "\\n",
                "print(f'Model created: {model}')"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "name": "python",
            "version": "3.11"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}

async def start_jupyter_server():
    """Start Jupyter server in the background."""
    import json
    
    {% if provision_nfs %}
    # Use persistent workspace
    workspace_dir = Path(WORKSPACE_PATH)
    {% else %}
    # Create temporary workspace
    workspace_dir = Path("/tmp/jupyter-workspace")
    {% endif %}
    workspace_dir.mkdir(exist_ok=True, parents=True)
    
    # Create a sample notebook
    sample_notebook = workspace_dir / "welcome.ipynb"
    if not sample_notebook.exists():
        sample_notebook.write_text(json.dumps(JUPYTER_NOTEBOOK_TEMPLATE, indent=2))
    
    # Generate Jupyter config
    jupyter_config = workspace_dir / "jupyter_config.py"
    jupyter_config.write_text("""
c.ServerApp.ip = '0.0.0.0'
c.ServerApp.port = 8888
c.ServerApp.allow_origin = '*'
c.ServerApp.allow_remote_access = True
c.ServerApp.disable_check_xsrf = True
c.ServerApp.token = ''
c.ServerApp.password = ''
c.ServerApp.root_dir = '{}'
""".format(str(workspace_dir)))
    
    # Start Jupyter server
    logger.info("Starting Jupyter server on port 8888")
    process = await asyncio.create_subprocess_exec(
        "jupyter", "lab",
        "--config", str(jupyter_config),
        "--no-browser",
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    
    # Give it time to start
    await asyncio.sleep(5)
    logger.info("Jupyter server started")
    
    return process

# Modal Function Configuration
@app.function(
    image=image,
    {% if gpu_type or deployment_mode == "gradio-jupyter" %}
    gpu="{{ gpu_type if gpu_type else 'any' }}",
    {% endif %}
    min_containers={{ min_containers }},
    max_containers={{ max_containers }},
    timeout={{ timeout_seconds }},
    memory=16384,  # 16GB RAM for Jupyter
    {% if provision_nfs %}
    volumes={WORKSPACE_PATH: volume},
    {% endif %}
)
@modal.concurrent(max_inputs=100)
@modal.asgi_app()
def deploy_gradio():
    """Deploy Gradio app with Jupyter notebooks 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 gradio-jupyter mode (GPU: {gpu_available})")
    
    # Start Jupyter in the background
    asyncio.create_task(start_jupyter_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 Jupyter integration
    with gr.Blocks() as enhanced_dashboard:
        gr.Markdown("# 🚀 Modal Deployment Dashboard")
        gr.Markdown("Monitor and manage your Modal deployment with Jupyter notebook integration")
        
        with gr.Tabs():
            # Jupyter Tab (first for easy access)
            with gr.Tab("🪐 Jupyter Notebooks"):
                gr.Markdown("### Interactive Jupyter Lab Environment")
                gr.Markdown("Access Jupyter Lab for data science and ML experiments.")
                
                with gr.Row():
                    gr.Markdown("**Jupyter is running on port 8888**")
                    jupyter_btn = gr.Button("🚀 Open Jupyter Lab", variant="primary", scale=2)
                
                gr.Markdown("""
                #### Features:
                - **JupyterLab interface** - Modern notebook environment
                - **GPU support** - Full CUDA acceleration
                - **Pre-installed packages** - All ML/AI libraries ready
                - **File browser** - Manage notebooks and data
                {% if provision_nfs %}
                - **Persistent storage** - Notebooks saved at `{{ WORKSPACE_PATH }}`
                {% else %}
                - **Temporary storage** - Notebooks for this session only
                {% endif %}
                
                #### Usage:
                1. Click the button above to open Jupyter Lab
                2. Create or open notebooks
                3. Run ML experiments with GPU acceleration
                """)
                
                # JavaScript to open Jupyter in new tab
                jupyter_btn.click(
                    None,
                    None,
                    None,
                    js="window.open('/jupyter', '_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) | 🪐 [Jupyter](/jupyter) | Generated by [modal-for-noobs](https://github.com/arthrod/modal-for-noobs)")
    
    enhanced_dashboard.queue(max_size=20)
    
    # FastAPI Setup with Jupyter proxy
    fastapi_app = FastAPI(
        title=APP_TITLE,
        description=APP_DESCRIPTION,
        version="1.0.0",
        docs_url="/docs",
        redoc_url="/redoc"
    )
    
    # Jupyter proxy client
    jupyter_client = httpx.AsyncClient(base_url="http://localhost:8888")
    
    # Add Jupyter proxy endpoints
    @fastapi_app.get("/jupyter")
    @fastapi_app.get("/jupyter/{path:path}")
    async def jupyter_proxy(request: Request, path: str = ""):
        """Proxy requests to Jupyter."""
        # 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 jupyter_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"Jupyter proxy error: {e}")
            return RedirectResponse(url="/")
    
    # WebSocket proxy for Jupyter
    @fastapi_app.websocket("/jupyter/{path:path}")
    async def jupyter_ws_proxy(websocket, path: str = ""):
        """Proxy WebSocket connections to Jupyter."""
        await websocket.accept()
        # This is a simplified version - full WS proxy would require more logic
        await websocket.close()
    
    # 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,
            "jupyter_enabled": True,
        }
    
    logger.info("Dashboard with Jupyter integration configured successfully")
    
    # Mount Gradio app
    return mount_gradio_app(fastapi_app, enhanced_dashboard, path="/")

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