sandbox.py

import base64
import time
import re
import mimetypes
import urllib.parse as _uparse
from typing import Dict, Optional

from code_processing import (
    is_streamlit_code, is_gradio_code, parse_multipage_html_output, 
    validate_and_autofix_files, inline_multipage_into_single_preview,
    build_transformers_inline_html
)

class SandboxRenderer:
    """Handles rendering of code in sandboxed environments"""
    
    @staticmethod
    def send_to_sandbox(code: str) -> str:
        """Render HTML in a sandboxed iframe"""
        html_doc = (code or "").strip()
        
        # Convert file:// URLs to data URIs for iframe compatibility
        html_doc = SandboxRenderer._inline_file_urls_as_data_uris(html_doc)
        
        # Encode and create iframe
        encoded_html = base64.b64encode(html_doc.encode('utf-8')).decode('utf-8')
        data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
        
        iframe = (
            f'<iframe src="{data_uri}" '
            f'width="100%" height="920px" '
            f'sandbox="allow-scripts allow-same-origin allow-forms allow-popups allow-modals allow-presentation" '
            f'allow="display-capture">'
            f'</iframe>'
        )
        
        return iframe
    
    @staticmethod
    def send_to_sandbox_with_refresh(code: str) -> str:
        """Render HTML with cache-busting for media generation updates"""
        html_doc = (code or "").strip()
        
        # Convert file:// URLs to data URIs for iframe compatibility
        html_doc = SandboxRenderer._inline_file_urls_as_data_uris(html_doc)
        
        # Add cache-busting timestamp
        timestamp = str(int(time.time() * 1000))
        cache_bust_comment = f"<!-- refresh-{timestamp} -->"
        html_doc = cache_bust_comment + html_doc
        
        # Encode and create iframe with unique key
        encoded_html = base64.b64encode(html_doc.encode('utf-8')).decode('utf-8')
        data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
        
        iframe = (
            f'<iframe src="{data_uri}" '
            f'width="100%" height="920px" '
            f'sandbox="allow-scripts allow-same-origin allow-forms allow-popups allow-modals allow-presentation" '
            f'allow="display-capture" '
            f'key="preview-{timestamp}">'
            f'</iframe>'
        )
        
        return iframe
    
    @staticmethod
    def _inline_file_urls_as_data_uris(html_content: str) -> str:
        """Convert file:// URLs to data URIs for iframe compatibility"""
        try:
            def _file_url_to_data_uri(file_url: str) -> Optional[str]:
                try:
                    parsed = _uparse.urlparse(file_url)
                    path = _uparse.unquote(parsed.path)
                    if not path:
                        return None
                    
                    with open(path, 'rb') as _f:
                        raw = _f.read()
                    
                    mime = mimetypes.guess_type(path)[0] or 'application/octet-stream'
                    b64 = base64.b64encode(raw).decode()
                    return f"data:{mime};base64,{b64}"
                except Exception:
                    return None
            
            def _repl_double(m):
                url = m.group(1)
                data_uri = _file_url_to_data_uri(url)
                return f'src="{data_uri}"' if data_uri else m.group(0)
            
            def _repl_single(m):
                url = m.group(1)
                data_uri = _file_url_to_data_uri(url)
                return f"src='{data_uri}'" if data_uri else m.group(0)
            
            # Replace file:// URLs in both single and double quotes
            html_content = re.sub(r'src="(file:[^"]+)"', _repl_double, html_content)
            html_content = re.sub(r"src='(file:[^']+)'", _repl_single, html_content)
            
        except Exception:
            # Best-effort; continue without inlining
            pass
        
        return html_content

class StreamlitRenderer:
    """Handles Streamlit app rendering using stlite"""
    
    @staticmethod
    def send_streamlit_to_stlite(code: str) -> str:
        """Render Streamlit code using stlite in sandboxed iframe"""
        html_doc = f"""<!doctype html>
<html>
  <head>
    <meta charset="UTF-8" />
    <meta http-equiv="X-UA-Compatible" content="IE=edge" />
    <meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=no" />
    <title>Streamlit Preview</title>
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/@stlite/[email protected]/build/stlite.css" />
    <style>
      html, body {{ margin: 0; padding: 0; height: 100%; }}
      streamlit-app {{ display: block; height: 100%; }}
      .stlite-loading {{
        display: flex;
        justify-content: center;
        align-items: center;
        height: 100vh;
        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
        color: #666;
      }}
    </style>
    <script type="module" src="https://cdn.jsdelivr.net/npm/@stlite/[email protected]/build/stlite.js"></script>
  </head>
  <body>
    <div class="stlite-loading">Loading Streamlit app...</div>
    <streamlit-app style="display: none;">
{code or ""}
    </streamlit-app>
    
    <script>
      // Show the app once stlite loads
      document.addEventListener('DOMContentLoaded', function() {{
        setTimeout(() => {{
          const loading = document.querySelector('.stlite-loading');
          const app = document.querySelector('streamlit-app');
          if (loading) loading.style.display = 'none';
          if (app) app.style.display = 'block';
        }}, 2000);
      }});
    </script>
  </body>
</html>"""
        
        encoded_html = base64.b64encode(html_doc.encode('utf-8')).decode('utf-8')
        data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
        
        iframe = (
            f'<iframe src="{data_uri}" '
            f'width="100%" height="920px" '
            f'sandbox="allow-scripts allow-same-origin allow-forms allow-popups allow-modals allow-presentation" '
            f'allow="display-capture">'
            f'</iframe>'
        )
        
        return iframe

class GradioRenderer:
    """Handles Gradio app rendering using gradio-lite"""
    
    @staticmethod
    def send_gradio_to_lite(code: str) -> str:
        """Render Gradio code using gradio-lite in sandboxed iframe"""
        html_doc = f"""<!doctype html>
<html>
  <head>
    <meta charset="UTF-8" />
    <meta http-equiv="X-UA-Compatible" content="IE=edge" />
    <meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=no" />
    <title>Gradio Preview</title>
    <script type="module" crossorigin src="https://cdn.jsdelivr.net/npm/@gradio/lite/dist/lite.js"></script>
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/@gradio/lite/dist/lite.css" />
    <style>
      html, body {{ margin: 0; padding: 0; height: 100%; }}
      gradio-lite {{ display: block; height: 100%; }}
      .gradio-loading {{
        display: flex;
        justify-content: center;
        align-items: center;
        height: 100vh;
        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
        color: #666;
        flex-direction: column;
        gap: 16px;
      }}
      .loading-spinner {{
        border: 3px solid #f3f3f3;
        border-top: 3px solid #ff6b6b;
        border-radius: 50%;
        width: 32px;
        height: 32px;
        animation: spin 1s linear infinite;
      }}
      @keyframes spin {{
        0% {{ transform: rotate(0deg); }}
        100% {{ transform: rotate(360deg); }}
      }}
    </style>
  </head>
  <body>
    <div class="gradio-loading">
      <div class="loading-spinner"></div>
      <div>Loading Gradio app...</div>
    </div>
    <gradio-lite style="display: none;">
{code or ""}
    </gradio-lite>
    
    <script>
      // Show the app once gradio-lite loads
      document.addEventListener('DOMContentLoaded', function() {{
        setTimeout(() => {{
          const loading = document.querySelector('.gradio-loading');
          const app = document.querySelector('gradio-lite');
          if (loading) loading.style.display = 'none';
          if (app) app.style.display = 'block';
        }}, 3000);
      }});
    </script>
  </body>
</html>"""
        
        encoded_html = base64.b64encode(html_doc.encode('utf-8')).decode('utf-8')
        data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
        
        iframe = (
            f'<iframe src="{data_uri}" '
            f'width="100%" height="920px" '
            f'sandbox="allow-scripts allow-same-origin allow-forms allow-popups allow-modals allow-presentation" '
            f'allow="display-capture">'
            f'</iframe>'
        )
        
        return iframe

class TransformersJSRenderer:
    """Handles Transformers.js app rendering"""
    
    @staticmethod
    def send_transformers_to_sandbox(files: Dict[str, str]) -> str:
        """Build and render transformers.js app in sandbox"""
        merged_html = build_transformers_inline_html(files)
        return SandboxRenderer.send_to_sandbox(merged_html)

class PreviewEngine:
    """Main preview engine that routes to appropriate renderers"""
    
    @staticmethod
    def generate_preview(code: str, language: str, **kwargs) -> str:
        """Generate appropriate preview based on code and language"""
        if not code or not code.strip():
            return PreviewEngine._create_empty_preview()
        
        try:
            if language == "html":
                return PreviewEngine._handle_html_preview(code)
            
            elif language == "streamlit" or (language == "python" and is_streamlit_code(code)):
                if is_streamlit_code(code):
                    return StreamlitRenderer.send_streamlit_to_stlite(code)
                else:
                    return PreviewEngine._create_info_preview(
                        "Streamlit Preview", 
                        "Add `import streamlit as st` to enable Streamlit preview."
                    )
            
            elif language == "gradio" or (language == "python" and is_gradio_code(code)):
                if is_gradio_code(code):
                    return GradioRenderer.send_gradio_to_lite(code)
                else:
                    return PreviewEngine._create_info_preview(
                        "Gradio Preview", 
                        "Add `import gradio as gr` to enable Gradio preview."
                    )
            
            elif language == "python":
                if is_streamlit_code(code):
                    return StreamlitRenderer.send_streamlit_to_stlite(code)
                elif is_gradio_code(code):
                    return GradioRenderer.send_gradio_to_lite(code)
                else:
                    return PreviewEngine._create_info_preview(
                        "Python Preview", 
                        "Preview available for Streamlit and Gradio apps. Add the appropriate import statements."
                    )
            
            elif language == "transformers.js":
                # Handle transformers.js with optional file parts
                html_part = kwargs.get('html_part', '')
                js_part = kwargs.get('js_part', '')
                css_part = kwargs.get('css_part', '')
                
                if html_part or js_part or css_part:
                    files = {'index.html': html_part or '', 'index.js': js_part or '', 'style.css': css_part or ''}
                else:
                    from code_processing import parse_transformers_js_output
                    files = parse_transformers_js_output(code)
                
                if files.get('index.html'):
                    return TransformersJSRenderer.send_transformers_to_sandbox(files)
                else:
                    return PreviewEngine._create_info_preview(
                        "Transformers.js Preview", 
                        "Generating transformers.js app... Please wait for all three files to be created."
                    )
            
            elif language == "svelte":
                return PreviewEngine._create_info_preview(
                    "Svelte Preview", 
                    "Preview is not available for Svelte apps. Download your code and deploy it to see the result."
                )
            
            else:
                return PreviewEngine._create_info_preview(
                    "Code Preview", 
                    f"Preview is not available for {language}. Supported: HTML, Streamlit, Gradio, Transformers.js"
                )
                
        except Exception as e:
            print(f"[Preview] Error generating preview: {str(e)}")
            return PreviewEngine._create_error_preview(f"Preview error: {str(e)}")
    
    @staticmethod
    def _handle_html_preview(code: str) -> str:
        """Handle HTML preview with multi-page support"""
        # Check for multi-page structure
        files = parse_multipage_html_output(code)
        files = validate_and_autofix_files(files)
        
        if files and files.get('index.html'):
            # Multi-page HTML
            merged = inline_multipage_into_single_preview(files)
            return SandboxRenderer.send_to_sandbox_with_refresh(merged)
        else:
            # Single HTML document
            from code_processing import extract_html_document
            safe_preview = extract_html_document(code)
            return SandboxRenderer.send_to_sandbox_with_refresh(safe_preview)
    
    @staticmethod
    def _create_empty_preview() -> str:
        """Create preview for empty content"""
        return PreviewEngine._create_info_preview("No Content", "Generate some code to see the preview.")
    
    @staticmethod
    def _create_info_preview(title: str, message: str) -> str:
        """Create informational preview"""
        return f"""<div style='padding: 2rem; text-align: center; color: #666; 
                   font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; 
                   background: #f8fafc; border-radius: 8px; margin: 1rem;'>
                   <h3 style='color: #374151; margin-top: 0;'>{title}</h3>
                   <p>{message}</p>
                   </div>"""
    
    @staticmethod
    def _create_error_preview(error_message: str) -> str:
        """Create error preview"""
        return f"""<div style='padding: 2rem; text-align: center; color: #dc2626; 
                   font-family: -apple-system, BlinkMacSystemFont, "Segeo UI", sans-serif; 
                   background: #fef2f2; border: 1px solid #fecaca; border-radius: 8px; margin: 1rem;'>
                   <h3 style='color: #dc2626; margin-top: 0;'>Preview Error</h3>
                   <p>{error_message}</p>
                   </div>"""

# Export main functions and classes
sandbox_renderer = SandboxRenderer()
streamlit_renderer = StreamlitRenderer()
gradio_renderer = GradioRenderer()
transformers_renderer = TransformersJSRenderer()
preview_engine = PreviewEngine()

# Main exports
def send_to_sandbox(code: str) -> str:
    return sandbox_renderer.send_to_sandbox(code)

def send_to_sandbox_with_refresh(code: str) -> str:
    return sandbox_renderer.send_to_sandbox_with_refresh(code)

def send_streamlit_to_stlite(code: str) -> str:
    return streamlit_renderer.send_streamlit_to_stlite(code)

def send_gradio_to_lite(code: str) -> str:
    return gradio_renderer.send_gradio_to_lite(code)

def send_transformers_to_sandbox(files: Dict[str, str]) -> str:
    return transformers_renderer.send_transformers_to_sandbox(files)

def generate_preview(code: str, language: str, **kwargs) -> str:
    return preview_engine.generate_preview(code, language, **kwargs)