import os
import subprocess
import time
import socket
import signal
import sys
from pyngrok import ngrok
import gradio as gr

# Get ngrok authtoken from environment
NGROK_AUTHTOKEN = os.getenv("NGROK_AUTHTOKEN")

# Global process references for cleanup
processes = {}

def is_port_open(port, host='localhost', timeout=2):
    """Check if a port is open and listening"""
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(timeout)
            result = sock.connect_ex((host, port))
            return result == 0
    except:
        return False

def cleanup_processes():
    """Cleanup all processes on exit"""
    print("\nCleaning up processes...")
    for name, proc in processes.items():
        try:
            if proc and proc.poll() is None:
                print(f"Terminating {name}...")
                proc.terminate()
                proc.wait(timeout=5)
        except:
            try:
                proc.kill()
            except:
                pass

def signal_handler(sig, frame):
    """Handle interrupt signals"""
    cleanup_processes()
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

def cleanup_previous_sessions():
    """Clean up any previous sessions"""
    print("Cleaning up previous sessions...")
    
    commands = [
        ["pkill", "-f", "x11vnc"],
        ["pkill", "-f", "websockify"],
        ["pkill", "-f", "Xvfb"],
        ["pkill", "-f", "fluxbox"],
        ["pkill", "-f", "organichits-exchanger"],
        ["fuser", "-k", "5901/tcp"],
        ["fuser", "-k", "8081/tcp"]
    ]
    
    for cmd in commands:
        try:
            subprocess.run(cmd, capture_output=True, timeout=5)
        except:
            pass
    
    time.sleep(2)

def setup_display():
    """Setup virtual display with proper settings for Electron/Chromium"""
    print("Setting up virtual display for Electron app...")
    
    # Kill any existing Xvfb on display 99
    subprocess.run(["pkill", "-f", "Xvfb :99"], capture_output=True)
    time.sleep(1)
    
    # Start Xvfb with settings optimized for Chromium
    xvfb_process = subprocess.Popen([
        "Xvfb", ":99",
        "-screen", "0", "1280x720x24",
        "-ac",
        "-nolisten", "tcp",
        "-dpi", "96",
        "+extension", "RANDR"
    ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    
    processes['xvfb'] = xvfb_process
    
    # Wait for X server to be ready
    time.sleep(3)
    
    # Set display environment globally
    os.environ['DISPLAY'] = ':99'
    
    # Verify display is working
    try:
        result = subprocess.run(
            ["xdpyinfo", "-display", ":99"],
            capture_output=True,
            timeout=5
        )
        if result.returncode == 0:
            print("✓ Virtual display :99 is ready")
        else:
            print("⚠ Display may not be fully ready")
    except:
        print("⚠ Could not verify display")
    
    # Start fluxbox window manager
    print("Starting window manager...")
    fluxbox_process = subprocess.Popen(
        ["fluxbox"],
        env={**os.environ, 'DISPLAY': ':99'},
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )
    
    processes['fluxbox'] = fluxbox_process
    time.sleep(2)
    
    return True

def start_vnc_server():
    """Start VNC server connected to virtual display"""
    print("Starting VNC server...")
    
    # Start x11vnc
    vnc_process = subprocess.Popen([
        "x11vnc",
        "-display", ":99",
        "-forever",
        "-shared",
        "-nopw",
        "-listen", "0.0.0.0",
        "-rfbport", "5901",
        "-xkb",
        "-noxrecord",
        "-noxfixes",
        "-noxdamage",
        "-wait", "10",
        "-defer", "10"
    ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    
    processes['vnc'] = vnc_process
    
    # Wait for VNC to start
    for i in range(15):
        if is_port_open(5901):
            print("✓ VNC server started on port 5901")
            return True
        time.sleep(1)
    
    print("✗ VNC server failed to start")
    return False

def start_organichits_app():
    """Start OrganicHits Exchanger with proper Electron/Chromium flags"""
    print("Starting OrganicHits Exchanger...")
    
    # Ensure display is set
    os.environ['DISPLAY'] = ':99'
    
    # Path to the binary and its directory
    app_dir = "./OrganicHits"
    app_path = os.path.join(app_dir, "organichits-exchanger")
    
    # Make sure it's executable
    os.chmod(app_path, 0o755)
    
    # Verify critical files exist
    critical_files = ['icudtl.dat', 'resources.pak', 'v8_context_snapshot.bin']
    for file in critical_files:
        file_path = os.path.join(app_dir, file)
        if not os.path.exists(file_path):
            print(f"⚠ Warning: Missing critical file: {file}")
    
    # Comprehensive Chromium flags for containerized Electron apps
    app_process = subprocess.Popen([
        app_path,
        "--no-sandbox",
        "--disable-setuid-sandbox",
        "--disable-dev-shm-usage",
        "--disable-gpu",
        "--disable-software-rasterizer",
        "--disable-gpu-compositing",
        "--enable-features=UseOzonePlatform",
        "--ozone-platform=x11",
        "--disable-features=VizDisplayCompositor",
        "--in-process-gpu",
        "--disable-accelerated-2d-canvas",
        "--disable-accelerated-video-decode",
        "--disable-accelerated-video-encode",
        "--no-first-run",
        "--no-default-browser-check",
        "--disable-background-timer-throttling",
        "--disable-backgrounding-occluded-windows",
        "--disable-renderer-backgrounding",
        "--disable-breakpad"
    ],
    env={
        **os.environ,
        'DISPLAY': ':99',
        'ELECTRON_DISABLE_SANDBOX': '1',
        'ELECTRON_ENABLE_LOGGING': '1',
        'LD_LIBRARY_PATH': app_dir
    },
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE,
    cwd=app_dir
    )
    
    processes['app'] = app_process
    
    # Give app time to start
    time.sleep(5)
    
    # Check if app is still running
    if app_process.poll() is None:
        print("✓ OrganicHits application started")
        return True
    else:
        stdout, stderr = app_process.communicate(timeout=1)
        print(f"✗ Application failed to start")
        print(f"stdout: {stdout.decode()[:500]}")
        print(f"stderr: {stderr.decode()[:500]}")
        return False

def start_websockify():
    """Start websockify proxy for web access"""
    print("Starting websockify on port 8081...")
    
    websockify_process = subprocess.Popen([
        "websockify",
        "8081",
        "localhost:5901",
        "--web=/usr/share/novnc/",
        "--verbose"
    ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    
    processes['websockify'] = websockify_process
    
    # Wait for websockify to start
    for i in range(10):
        if is_port_open(8081):
            print("✓ Websockify started on port 8081")
            return True
        time.sleep(1)
    
    print("✗ Websockify failed to start")
    return False

def start_vnc_and_ngrok():
    """Main function to start everything"""
    if not NGROK_AUTHTOKEN:
        return "❌ Error: NGROK_AUTHTOKEN not found in environment secrets!"

    try:
        # Step 1: Cleanup
        cleanup_previous_sessions()
        
        # Step 2: Setup display
        if not setup_display():
            return "❌ Failed to setup virtual display"
        
        # Step 3: Start VNC
        if not start_vnc_server():
            return "❌ Failed to start VNC server"
        
        # Step 4: Start websockify
        if not start_websockify():
            return "❌ Failed to start websockify"
        
        # Step 5: Start the application
        app_started = start_organichits_app()
        app_status = "✓ Running" if app_started else "⚠ May have failed (check logs)"
        
        # Step 6: Setup ngrok
        print("Creating ngrok tunnel...")
        try:
            ngrok.set_auth_token(NGROK_AUTHTOKEN)
            public_url = ngrok.connect(8081, "http")
            novnc_url = f"{public_url}/vnc.html?autoconnect=true"
            
            return f"""✅ VNC Session Started Successfully!

🌐 Access Your Application:
{novnc_url}

📊 Status:
• Virtual Display: ✓ Running on :99
• VNC Server: ✓ Running on port 5901
• Websockify: ✓ Running on port 8081
• OrganicHits App: {app_status}

📱 Click the URL above to access your OrganicHits Exchanger!

⚠️ Notes:
- First load may take 10-20 seconds
- Click "Connect" if not auto-connected
- No password required
- Keep this tab open to maintain the session

🔧 Troubleshooting:
- If black screen persists, wait 30 seconds and refresh
- Check "Service Status" below
- Application logs are being captured
"""
            
        except Exception as e:
            return f"❌ Failed to create ngrok tunnel: {e}"

    except Exception as e:
        return f"❌ Unexpected error: {str(e)}"

def check_services():
    """Check status of all services"""
    services = {
        'Virtual Display (Xvfb)': is_port_open(6099) or ('xvfb' in processes and processes['xvfb'].poll() is None),
        'VNC Server': is_port_open(5901),
        'Websockify': is_port_open(8081),
        'OrganicHits App': 'app' in processes and processes['app'].poll() is None
    }
    
    status_text = "🔍 Service Status:\n\n"
    for service, running in services.items():
        emoji = "✅" if running else "❌"
        status_text += f"{emoji} {service}\n"
    
    return status_text

def get_app_logs():
    """Get application logs"""
    if 'app' not in processes or processes['app'].poll() is not None:
        return "Application not running or has stopped"
    
    try:
        # Try to get recent output
        return "Application is running. Check console for detailed logs."
    except:
        return "Could not retrieve logs"

def main():
    with gr.Blocks(title="OrganicHits VNC Remote Desktop", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🖥️ OrganicHits Exchanger - Remote Desktop
        
        Run OrganicHits Exchanger in your browser via VNC
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                output = gr.Textbox(
                    label="🌐 Access Information",
                    lines=15,
                    placeholder="Click 'Start VNC Session' to begin...",
                    show_copy_button=True
                )
                
                with gr.Row():
                    start_button = gr.Button(
                        "🚀 Start VNC Session",
                        variant="primary",
                        size="lg",
                        scale=2
                    )
                    check_btn = gr.Button(
                        "🔍 Check Status",
                        variant="secondary",
                        scale=1
                    )
            
            with gr.Column(scale=1):
                status = gr.Textbox(
                    label="📊 Service Status",
                    lines=8,
                    value="Services not started yet..."
                )
                
                gr.Markdown("""
                ### ⏱️ Expected Timeline:
                - Display setup: 3-5s
                - VNC start: 2-3s
                - App launch: 5-10s
                - **Total: ~20-30 seconds**
                """)
        
        gr.Markdown("""
        ---
        ### 📋 How to Use:
        
        1. **Click "Start VNC Session"** and wait ~30 seconds
        2. **Click the URL** that appears above
        3. **Wait for connection** - you'll see the desktop
        4. **OrganicHits will auto-launch** in the window
        
        ### 🎯 What You'll See:
        - A desktop environment (Fluxbox)
        - OrganicHits Exchanger window
        - You can interact with it normally!
        
        ### ⚠️ Important Notes:
        - Keep this browser tab open
        - No authentication (temporary session)
        - Session ends if you close this page
        - First connection may show black screen for 10-20 seconds (normal!)
        
        ### 🔧 Troubleshooting:
        - **Black screen?** Wait 30 seconds and refresh the VNC page
        - **Can't connect?** Check if NGROK_AUTHTOKEN is set in Secrets
        - **App not showing?** Click "Check Status" to diagnose
        """)
        
        start_button.click(
            fn=start_vnc_and_ngrok,
            outputs=output
        )
        
        check_btn.click(
            fn=check_services,
            outputs=status
        )

    # Launch Gradio
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )

if __name__ == "__main__":
    main()