import pygame
import numpy as np
from flask import Flask, Response, render_template_string, request, jsonify
import time
import os
import cv2
import subprocess
import threading
import queue
import uuid

# Initialize Pygame headlessly
os.environ['SDL_VIDEODRIVER'] = 'dummy'
pygame.init()

# Try to initialize mixer, but don't crash if it fails
try:
    pygame.mixer.init(frequency=44100, size=-16, channels=2)
    print("✅ Audio mixer initialized")
except Exception as e:
    print(f"⚠️ Audio mixer not available: {e}")
    print("   Continuing without sound...")

app = Flask(__name__)

class ShaderRenderer:
    def __init__(self, width=640, height=480):
        self.width = width
        self.height = height
        self.mouse_x = width // 2
        self.mouse_y = height // 2
        self.start_time = time.time()
        self.surface = pygame.Surface((width, height))
        self.frame_count = 0
        self.last_frame_time = time.time()
        self.fps = 0
        
        # Sound sources
        self.sound_source = 'none'
        self.pygame_sound = None
        self.pygame_playing = False
        self.sound_amp = 0.0
        
        # Load pygame sound if available
        if os.path.exists('sound.mp3'):
            try:
                self.pygame_sound = pygame.mixer.Sound('sound.mp3')
                print("✅ Pygame sound loaded")
            except:
                print("⚠️ Could not load sound.mp3")
    
    def set_mouse(self, x, y):
        self.mouse_x = max(0, min(self.width, x))
        self.mouse_y = max(0, min(self.height, y))
    
    def set_sound_source(self, source):
        """Change sound source: none, pygame, browser"""
        self.sound_source = source
        
        # Handle pygame sound
        if source == 'pygame':
            if self.pygame_sound and not self.pygame_playing:
                self.pygame_sound.play(loops=-1)
                self.pygame_playing = True
                self.sound_amp = 0.5
        else:
            if self.pygame_playing:
                pygame.mixer.stop()
                self.pygame_playing = False
                self.sound_amp = 0.0
    
    def render_frame(self):
        """Render the pygame frame"""
        t = time.time() - self.start_time
        
        # Calculate FPS
        self.frame_count += 1
        if time.time() - self.last_frame_time > 1.0:
            self.fps = self.frame_count
            self.frame_count = 0
            self.last_frame_time = time.time()
        
        # Clear
        self.surface.fill((20, 20, 30))
        
        # Draw TOP marker
        pygame.draw.rect(self.surface, (255, 100, 100), (10, 10, 100, 30))
        font = pygame.font.Font(None, 24)
        text = font.render("TOP", True, (255, 255, 255))
        self.surface.blit(text, (20, 15))
        
        # Draw BOTTOM marker
        pygame.draw.rect(self.surface, (100, 255, 100), (10, self.height-40, 100, 30))
        text = font.render("BOTTOM", True, (0, 0, 0))
        self.surface.blit(text, (20, self.height-35))
        
        # Draw moving circle
        circle_size = 30 + int(20 * np.sin(t * 2))
        
        # Color code based on sound source
        if self.sound_source == 'pygame':
            color = (100, 255, 100)  # Green for pygame sound
        elif self.sound_source == 'browser':
            color = (100, 100, 255)  # Blue for browser sound
        else:
            color = (255, 100, 100)  # Red for no sound
            
        pygame.draw.circle(self.surface, color, 
                         (self.mouse_x, self.mouse_y), circle_size)
        
        # Draw grid
        for x in range(0, self.width, 50):
            alpha = int(40 + 20 * np.sin(x * 0.1 + t))
            pygame.draw.line(self.surface, (alpha, alpha, 50), 
                           (x, 0), (x, self.height))
        for y in range(0, self.height, 50):
            alpha = int(40 + 20 * np.cos(y * 0.1 + t))
            pygame.draw.line(self.surface, (alpha, alpha, 50), 
                           (0, y), (self.width, y))
        
        # Sound meter
        meter_width = int(200 * self.sound_amp)
        pygame.draw.rect(self.surface, (60, 60, 60), (self.width-220, 10, 200, 20))
        pygame.draw.rect(self.surface, (100, 255, 100), 
                        (self.width-220, 10, meter_width, 20))
        
        # FPS counter
        fps_text = font.render(f"FPS: {self.fps}", True, (255, 255, 0))
        self.surface.blit(fps_text, (self.width-150, self.height-60))
        
        return pygame.image.tostring(self.surface, 'RGB')
    
    def get_frame(self):
        return self.render_frame()
    
    def get_frame_jpeg(self, quality=80):
        """Return frame as JPEG"""
        frame = self.get_frame()
        # Convert to numpy array for OpenCV
        img = np.frombuffer(frame, dtype=np.uint8).reshape((self.height, self.width, 3))
        # Convert RGB to BGR for OpenCV
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        # Encode as JPEG
        _, jpeg = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, quality])
        return jpeg.tobytes()

renderer = ShaderRenderer()

# Streaming managers
class StreamManager:
    def __init__(self):
        self.streams = {}
        
    def create_mjpeg_stream(self):
        stream_id = str(uuid.uuid4())
        self.streams[stream_id] = {
            'type': 'mjpeg',
            'active': True,
            'clients': 0
        }
        return stream_id
    
    def create_ffmpeg_stream(self, format_type='webm'):
        stream_id = str(uuid.uuid4())
        
        # FFmpeg command based on format
        if format_type == 'webm':
            cmd = [
                'ffmpeg',
                '-f', 'rawvideo',
                '-pix_fmt', 'rgb24',
                '-s', '640x480',
                '-r', '30',
                '-i', '-',
                '-c:v', 'libvpx-vp9',
                '-b:v', '1M',
                '-cpu-used', '4',
                '-deadline', 'realtime',
                '-f', 'webm',
                '-'
            ]
            mimetype = 'video/webm'
        else:  # mp4
            cmd = [
                'ffmpeg',
                '-f', 'rawvideo',
                '-pix_fmt', 'rgb24',
                '-s', '640x480',
                '-r', '30',
                '-i', '-',
                '-c:v', 'libx264',
                '-preset', 'ultrafast',
                '-tune', 'zerolatency',
                '-b:v', '1M',
                '-f', 'mp4',
                '-movflags', 'frag_keyframe+empty_moov',
                '-'
            ]
            mimetype = 'video/mp4'
        
        # Start FFmpeg process
        process = subprocess.Popen(
            cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            bufsize=0
        )
        
        frame_queue = queue.Queue(maxsize=30)
        
        self.streams[stream_id] = {
            'type': format_type,
            'mimetype': mimetype,
            'active': True,
            'process': process,
            'queue': frame_queue,
            'clients': 0
        }
        
        # Start frame pusher thread
        def push_frames():
            while self.streams.get(stream_id, {}).get('active', False):
                try:
                    frame = renderer.get_frame()
                    process.stdin.write(frame)
                except:
                    break
            process.terminate()
        
        threading.Thread(target=push_frames, daemon=True).start()
        
        return stream_id
    
    def get_stream(self, stream_id):
        return self.streams.get(stream_id)
    
    def close_stream(self, stream_id):
        if stream_id in self.streams:
            stream = self.streams[stream_id]
            if 'process' in stream:
                stream['process'].terminate()
            del self.streams[stream_id]

stream_manager = StreamManager()

@app.route('/')
def index():
    return render_template_string('''
    <!DOCTYPE html>
    <html>
    <head>
        <title>🎮 Pygame + Multi-format Streaming</title>
        <style>
            body {
                margin: 0;
                background: #0a0a0a;
                color: white;
                font-family: 'Segoe UI', Arial, sans-serif;
                display: flex;
                justify-content: center;
                align-items: center;
                min-height: 100vh;
            }
            
            .container {
                max-width: 900px;
                padding: 20px;
                text-align: center;
            }
            
            h1 {
                color: #4CAF50;
                margin-bottom: 20px;
                text-shadow: 0 0 10px rgba(76, 175, 80, 0.3);
            }
            
            .video-container {
                background: #000;
                border-radius: 12px;
                padding: 5px;
                margin: 20px 0;
                box-shadow: 0 0 30px rgba(76, 175, 80, 0.2);
            }
            
            #videoPlayer {
                width: 100%;
                max-width: 640px;
                height: auto;
                border-radius: 8px;
                display: block;
                margin: 0 auto;
                background: #111;
            }
            
            #mjpegImg {
                width: 100%;
                max-width: 640px;
                height: auto;
                border-radius: 8px;
                display: none;
                margin: 0 auto;
                background: #111;
            }
            
            .controls {
                background: #1a1a1a;
                border-radius: 12px;
                padding: 20px;
                margin-top: 20px;
            }
            
            .format-buttons {
                display: flex;
                gap: 10px;
                justify-content: center;
                margin: 20px 0;
                flex-wrap: wrap;
            }
            
            button {
                background: #333;
                color: white;
                border: none;
                padding: 12px 24px;
                font-size: 16px;
                border-radius: 8px;
                cursor: pointer;
                transition: all 0.3s;
                min-width: 100px;
                font-weight: bold;
                border: 1px solid #444;
            }
            
            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 5px 15px rgba(0,0,0,0.3);
            }
            
            button.active {
                background: #4CAF50;
                border-color: #4CAF50;
                box-shadow: 0 0 20px #4CAF50;
            }
            
            .status-panel {
                background: #222;
                border-radius: 8px;
                padding: 15px;
                margin-top: 20px;
                display: flex;
                justify-content: space-around;
                align-items: center;
                flex-wrap: wrap;
                gap: 15px;
            }
            
            .status-item {
                display: flex;
                align-items: center;
                gap: 10px;
            }
            
            .status-label {
                color: #888;
                font-size: 14px;
            }
            
            .status-value {
                background: #333;
                padding: 5px 12px;
                border-radius: 20px;
                font-size: 14px;
                font-weight: bold;
            }
            
            .browser-support {
                font-size: 12px;
                color: #666;
                margin-top: 15px;
                padding: 10px;
                border-top: 1px solid #333;
            }
            
            .badge {
                display: inline-block;
                padding: 3px 8px;
                border-radius: 4px;
                font-size: 11px;
                margin-left: 5px;
            }
            
            .badge.green {
                background: #4CAF50;
            }
            
            .badge.yellow {
                background: #ffaa00;
                color: black;
            }
            
            .badge.red {
                background: #ff4444;
            }
            
            .info-text {
                color: #666;
                font-size: 12px;
                margin-top: 10px;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎮 Pygame + Multi-format Streaming</h1>
            
            <div class="video-container">
                <video id="videoPlayer" autoplay controls muted></video>
                <img id="mjpegImg" src="">
            </div>
            
            <div class="controls">
                <h3>📡 Streaming Format</h3>
                <div class="format-buttons">
                    <button id="btnMjpeg" onclick="setFormat('mjpeg')" class="active">📸 MJPEG</button>
                    <button id="btnWebm" onclick="setFormat('webm')">🎥 WebM (VP9)</button>
                    <button id="btnMp4" onclick="setFormat('mp4')">🎬 MP4 (H.264)</button>
                </div>
                
                <div class="status-panel">
                    <div class="status-item">
                        <span class="status-label">Current Format:</span>
                        <span id="currentFormat" class="status-value">MJPEG</span>
                    </div>
                    <div class="status-item">
                        <span class="status-label">Browser Support:</span>
                        <span id="browserSupport" class="status-value">Checking...</span>
                    </div>
                    <div class="status-item">
                        <span class="status-label">Stream Status:</span>
                        <span id="streamStatus" class="status-value">🟢 Active</span>
                    </div>
                </div>
                
                <div class="browser-support" id="supportDetails">
                    Testing browser capabilities...
                </div>
                
                <div class="info-text">
                    ⚡ MJPEG: Lowest CPU, universal support<br>
                    🎥 WebM: Efficient compression, best for Chrome/Firefox<br>
                    🎬 MP4: Universal support, hardware accelerated
                </div>
            </div>
        </div>
        
        <script>
            const videoPlayer = document.getElementById('videoPlayer');
            const mjpegImg = document.getElementById('mjpegImg');
            let currentFormat = 'mjpeg';
            let streamActive = true;
            
            // Check browser capabilities
            function checkBrowserSupport() {
                const video = document.createElement('video');
                const support = {
                    webm: video.canPlayType('video/webm; codecs="vp9, vorbis"'),
                    webmVP8: video.canPlayType('video/webm; codecs="vp8, vorbis"'),
                    mp4: video.canPlayType('video/mp4; codecs="avc1.42E01E, mp4a.40.2"'),
                    mp4H264: video.canPlayType('video/mp4; codecs="avc1.64001E"'),
                    mjpeg: true
                };
                
                let supportText = '';
                let details = '';
                
                if (support.webm) {
                    supportText += '✓ WebM VP9 ';
                    details += 'WebM VP9: ' + support.webm + '<br>';
                }
                if (support.mp4) {
                    supportText += '✓ MP4 H.264 ';
                    details += 'MP4 H.264: ' + support.mp4 + '<br>';
                }
                supportText += '✓ MJPEG';
                
                document.getElementById('browserSupport').innerHTML = supportText;
                document.getElementById('supportDetails').innerHTML = details;
                
                return support;
            }
            
            function setFormat(format) {
                currentFormat = format;
                
                // Update button states
                document.getElementById('btnMjpeg').className = format === 'mjpeg' ? 'active' : '';
                document.getElementById('btnWebm').className = format === 'webm' ? 'active' : '';
                document.getElementById('btnMp4').className = format === 'mp4' ? 'active' : '';
                
                document.getElementById('currentFormat').innerHTML = format.toUpperCase();
                
                // Hide both players first
                videoPlayer.style.display = 'none';
                mjpegImg.style.display = 'none';
                
                if (format === 'mjpeg') {
                    // Use img tag for MJPEG
                    mjpegImg.style.display = 'block';
                    mjpegImg.src = '/video/mjpeg?' + Date.now(); // Add timestamp to prevent caching
                } else {
                    // Use video tag for WebM/MP4
                    videoPlayer.style.display = 'block';
                    
                    // Stop current playback
                    videoPlayer.pause();
                    videoPlayer.removeAttribute('src');
                    videoPlayer.load();
                    
                    // Set new source
                    const source = document.createElement('source');
                    if (format === 'webm') {
                        source.src = '/video/webm?' + Date.now();
                        source.type = 'video/webm';
                    } else {
                        source.src = '/video/mp4?' + Date.now();
                        source.type = 'video/mp4';
                    }
                    
                    videoPlayer.appendChild(source);
                    videoPlayer.load();
                    videoPlayer.play().catch(e => console.log('Playback error:', e));
                }
            }
            
            // Monitor stream health
            function checkStreamHealth() {
                if (currentFormat === 'mjpeg') {
                    // For MJPEG, check if image is loading
                    mjpegImg.onerror = function() {
                        document.getElementById('streamStatus').innerHTML = '🔴 Error';
                        streamActive = false;
                    };
                    mjpegImg.onload = function() {
                        document.getElementById('streamStatus').innerHTML = '🟢 Active';
                        streamActive = true;
                    };
                } else {
                    // For video, check if playing
                    videoPlayer.onerror = function() {
                        document.getElementById('streamStatus').innerHTML = '🔴 Error';
                        streamActive = false;
                    };
                    videoPlayer.onplaying = function() {
                        document.getElementById('streamStatus').innerHTML = '🟢 Active';
                        streamActive = true;
                    };
                }
            }
            
            // Initialize
            checkBrowserSupport();
            setFormat('mjpeg');  // Start with MJPEG for reliability
            checkStreamHealth();
            
            // Reconnect on error
            setInterval(() => {
                if (!streamActive) {
                    console.log('Attempting to reconnect...');
                    setFormat(currentFormat);
                }
            }, 5000);
        </script>
    </body>
    </html>
    ''')

@app.route('/video/mjpeg')
def video_mjpeg():
    """MJPEG streaming endpoint"""
    def generate():
        while True:
            frame = renderer.get_frame_jpeg(quality=70)
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
            # Small delay to control frame rate
            time.sleep(1/30)
    
    return Response(
        generate(),
        mimetype='multipart/x-mixed-replace; boundary=frame'
    )

@app.route('/video/webm')
def video_webm():
    """WebM streaming endpoint"""
    cmd = [
        'ffmpeg',
        '-f', 'rawvideo',
        '-pix_fmt', 'rgb24',
        '-s', '640x480',
        '-r', '30',
        '-i', '-',
        '-c:v', 'libvpx-vp9',
        '-b:v', '1M',
        '-cpu-used', '4',
        '-deadline', 'realtime',
        '-f', 'webm',
        '-'
    ]
    
    process = subprocess.Popen(
        cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        bufsize=0
    )
    
    def generate():
        # Frame pusher thread
        def push_frames():
            while True:
                try:
                    frame = renderer.get_frame()
                    process.stdin.write(frame)
                except:
                    break
        
        threading.Thread(target=push_frames, daemon=True).start()
        
        # Read and yield encoded data
        while True:
            data = process.stdout.read(4096)
            if not data:
                break
            yield data
    
    return Response(
        generate(),
        mimetype='video/webm',
        headers={
            'Cache-Control': 'no-cache',
            'Transfer-Encoding': 'chunked'
        }
    )

@app.route('/video/mp4')
def video_mp4():
    """MP4 streaming endpoint"""
    cmd = [
        'ffmpeg',
        '-f', 'rawvideo',
        '-pix_fmt', 'rgb24',
        '-s', '640x480',
        '-r', '30',
        '-i', '-',
        '-c:v', 'libx264',
        '-preset', 'ultrafast',
        '-tune', 'zerolatency',
        '-b:v', '1M',
        '-f', 'mp4',
        '-movflags', 'frag_keyframe+empty_moov',
        '-'
    ]
    
    process = subprocess.Popen(
        cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        bufsize=0
    )
    
    def generate():
        def push_frames():
            while True:
                try:
                    frame = renderer.get_frame()
                    process.stdin.write(frame)
                except:
                    break
        
        threading.Thread(target=push_frames, daemon=True).start()
        
        while True:
            data = process.stdout.read(4096)
            if not data:
                break
            yield data
    
    return Response(
        generate(),
        mimetype='video/mp4',
        headers={
            'Cache-Control': 'no-cache',
            'Transfer-Encoding': 'chunked'
        }
    )

@app.route('/mouse', methods=['POST'])
def mouse():
    """Keep mouse endpoint for interactivity"""
    data = request.json
    renderer.set_mouse(data['x'], data['y'])
    return 'OK'

@app.route('/sound/source', methods=['POST'])
def sound_source():
    """Keep sound source endpoint"""
    data = request.json
    renderer.set_sound_source(data['source'])
    return 'OK'

@app.route('/sound/amp')
def sound_amp():
    """Keep sound amp endpoint for compatibility"""
    return {'amp': renderer.sound_amp}

@app.route('/static/sound.mp3')
def serve_sound():
    """Serve sound file"""
    if os.path.exists('sound.mp3'):
        with open('sound.mp3', 'rb') as f:
            return Response(f.read(), mimetype='audio/mpeg')
    return 'Sound not found', 404

if __name__ == '__main__':
    print("\n" + "="*70)
    print("🎮 Pygame + Multi-format Streaming")
    print("="*70)
    print("📡 Streaming endpoints:")
    print("   • MJPEG: /video/mjpeg - Low CPU, universal")
    print("   • WebM:  /video/webm  - VP9 codec, efficient")
    print("   • MP4:   /video/mp4   - H.264, hardware accelerated")
    print("\n🌐 Main page: /")
    print("="*70 + "\n")
    
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)