import pygame
import numpy as np
from flask import Flask, Response, render_template_string, request
import time
import os

# Initialize Pygame headlessly
os.environ['SDL_VIDEODRIVER'] = 'dummy'
pygame.init()
pygame.mixer.init(frequency=44100, size=-16, channels=2)

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))
        
        # Sound sources
        self.sound_source = 'none'  # 'none', 'pygame', 'browser'
        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 a simple pattern that shader will transform"""
        t = time.time() - self.start_time
        
        # Clear
        self.surface.fill((20, 20, 30))
        
        # Draw TOP marker (should be at top of canvas)
        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))
        
        # Shader indicator
        shader_text = font.render("SHADER ON", True, (255, 255, 0))
        self.surface.blit(shader_text, (self.width-150, self.height-30))
        
        return pygame.image.tostring(self.surface, 'RGB')
    
    def get_frame(self):
        return self.render_frame()

renderer = ShaderRenderer()

@app.route('/')
def index():
    return render_template_string('''
    <!DOCTYPE html>
    <html>
    <head>
        <title>Pygame + WebGL Shader + Sound</title>
        <style>
            body {
                margin: 0;
                background: #0a0a0a;
                color: white;
                font-family: Arial, sans-serif;
                display: flex;
                flex-direction: column;
                justify-content: center;
                align-items: center;
                min-height: 100vh;
            }
            
            .container {
                text-align: center;
            }
            
            canvas {
                width: 640px;
                height: 480px;
                border: 3px solid #4CAF50;
                border-radius: 8px;
                cursor: crosshair;
                display: block;
                margin: 20px 0;
            }
            
            .controls {
                background: #1a1a1a;
                padding: 20px;
                border-radius: 8px;
                margin-top: 20px;
            }
            
            .button-group {
                display: flex;
                gap: 10px;
                justify-content: center;
                margin: 15px 0;
                flex-wrap: wrap;
            }
            
            button {
                background: #333;
                color: white;
                border: none;
                padding: 12px 24px;
                font-size: 16px;
                border-radius: 5px;
                cursor: pointer;
                transition: all 0.3s;
                min-width: 120px;
            }
            
            button:hover {
                transform: scale(1.05);
            }
            
            button.active {
                background: #4CAF50;
                box-shadow: 0 0 20px #4CAF50;
            }
            
            button.shader.active {
                background: #ffaa00;
                color: black;
            }
            
            button.sound.active {
                background: #4CAF50;
            }
            
            button.browser.active {
                background: #2196F3;
            }
            
            .status-bar {
                display: flex;
                justify-content: space-between;
                align-items: center;
                margin-top: 15px;
                padding: 10px;
                background: #222;
                border-radius: 5px;
            }
            
            .indicator {
                padding: 5px 10px;
                border-radius: 4px;
                font-size: 14px;
            }
            
            .indicator.none { background: #444; }
            .indicator.pygame { background: #4CAF50; }
            .indicator.browser { background: #2196F3; }
            .indicator.shader-on { background: #ffaa00; color: black; }
            .indicator.shader-off { background: #666; }
            
            .meter {
                width: 200px;
                height: 20px;
                background: #333;
                border-radius: 10px;
                overflow: hidden;
            }
            
            .meter-fill {
                height: 100%;
                width: 0%;
                background: linear-gradient(90deg, #4CAF50, #2196F3);
                transition: width 0.05s;
            }
            
            .badge {
                display: inline-block;
                padding: 5px 10px;
                border-radius: 4px;
                margin-left: 10px;
                font-size: 12px;
            }
            
            .badge.on {
                background: #4CAF50;
                color: white;
            }
            
            .badge.off {
                background: #ff4444;
                color: white;
            }
            
            .color-sample {
                display: inline-block;
                width: 12px;
                height: 12px;
                border-radius: 3px;
                margin: 0 5px;
            }
            
            .info-text {
                font-size: 12px;
                color: #666;
                margin-top: 15px;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎮 Pygame + WebGL Shader + Sound</h1>
            
            <canvas id="canvas" width="640" height="480"></canvas>
            
            <div class="controls">
                <div class="button-group">
                    <button id="shaderBtn" class="shader active" onclick="toggleShader()">
                        🔮 SHADER ON
                    </button>
                    <span id="shaderBadge" class="badge on">EFFECTS ACTIVE</span>
                </div>
                
                <h3>🔊 Sound Source</h3>
                <div class="button-group">
                    <button id="btnNone" class="sound active" onclick="setSoundSource('none')">🔇 None</button>
                    <button id="btnPygame" class="sound" onclick="setSoundSource('pygame')">🎮 Pygame</button>
                    <button id="btnBrowser" class="sound" onclick="setSoundSource('browser')">🌐 Browser</button>
                </div>
                
                <div class="status-bar">
                    <div>
                        <span id="sourceIndicator" class="indicator none">Source: None</span>
                        <span id="shaderIndicator" class="indicator shader-on" style="margin-left: 10px;">Shader: ON</span>
                    </div>
                    <div class="meter">
                        <div id="soundMeter" class="meter-fill"></div>
                    </div>
                </div>
                
                <div class="info-text">
                    <span class="color-sample" style="background: #ff6464;"></span> No sound
                    <span class="color-sample" style="background: #64ff64;"></span> Pygame sound
                    <span class="color-sample" style="background: #6464ff;"></span> Browser sound
                    | <span style="color: #ff6464;">🔴 TOP</span> marker should be at top
                    | <span style="color: #64ff64;">🟢 BOTTOM</span> at bottom
                </div>
            </div>
        </div>
        
        <!-- Audio element for browser sound -->
        <audio id="browserAudio" loop style="display:none;">
            <source src="/static/sound.mp3" type="audio/mpeg">
        </audio>
        
        <!-- WebGL Shader -->
        <script id="vertex-shader" type="x-shader/x-vertex">
            attribute vec2 position;
            varying vec2 vUv;
            void main() {
                // Flip Y coordinate to fix Pygame orientation
                vUv = vec2(position.x * 0.5 + 0.5, 1.0 - (position.y * 0.5 + 0.5));
                gl_Position = vec4(position, 0.0, 1.0);
            }
        </script>
        
        <script id="fragment-shader" type="x-shader/x-fragment">
            precision highp float;
            
            uniform sampler2D uTexture;
            uniform float uTime;
            uniform vec2 uMouse;
            uniform vec2 uResolution;
            uniform bool uShaderEnabled;
            
            varying vec2 vUv;
            
            void main() {
                // Get pixel from pygame texture (now correctly oriented)
                vec4 color = texture2D(uTexture, vUv);
                
                if (uShaderEnabled) {
                    // SHADER EFFECTS ON
                    
                    // 1. Time-based color shift
                    color.r += sin(uTime + vUv.x * 10.0) * 0.2;
                    color.g += cos(uTime + vUv.y * 10.0) * 0.2;
                    
                    // 2. Mouse ripple effect
                    float dist = distance(vUv, uMouse);
                    if (dist < 0.2) {
                        float ripple = sin(dist * 50.0 - uTime * 5.0) * 0.5 + 0.5;
                        color.rgb += vec3(0.5, 0.2, 0.8) * ripple * 0.5;
                    }
                    
                    // 3. Scanlines
                    float scanline = sin(vUv.y * uResolution.y * 2.0 + uTime * 10.0) * 0.1 + 0.9;
                    color.rgb *= scanline;
                    
                    // 4. Edge glow
                    float edge = 1.0 - abs(vUv.x - 0.5) * 2.0;
                    edge *= 1.0 - abs(vUv.y - 0.5) * 2.0;
                    color.rgb += vec3(0.2, 0.1, 0.5) * edge * sin(uTime) * 0.3;
                }
                // else: SHADER EFFECTS OFF - pure Pygame pixels
                
                gl_FragColor = color;
            }
        </script>
        
        <script>
            const canvas = document.getElementById('canvas');
            const gl = canvas.getContext('webgl');
            const browserAudio = document.getElementById('browserAudio');
            
            if (!gl) {
                alert('WebGL not supported!');
            }
            
            let texture = gl.createTexture();
            let startTime = Date.now() / 1000;
            let shaderEnabled = true;
            let mouse = [0.5, 0.5];
            let currentSource = 'none';
            
            // UI Elements
            const shaderBtn = document.getElementById('shaderBtn');
            const shaderBadge = document.getElementById('shaderBadge');
            const shaderIndicator = document.getElementById('shaderIndicator');
            
            function toggleShader() {
                shaderEnabled = !shaderEnabled;
                
                // Update UI
                if (shaderEnabled) {
                    shaderBtn.className = 'shader active';
                    shaderBtn.innerHTML = '🔮 SHADER ON';
                    shaderBadge.className = 'badge on';
                    shaderBadge.innerHTML = 'EFFECTS ACTIVE';
                    shaderIndicator.className = 'indicator shader-on';
                    shaderIndicator.innerHTML = 'Shader: ON';
                } else {
                    shaderBtn.className = '';
                    shaderBtn.innerHTML = '🎮 SHADER OFF';
                    shaderBadge.className = 'badge off';
                    shaderBadge.innerHTML = 'PURE PYGAME';
                    shaderIndicator.className = 'indicator shader-off';
                    shaderIndicator.innerHTML = 'Shader: OFF';
                }
                
                // Update shader uniform
                if (program) {
                    gl.useProgram(program);
                    const uShaderEnabled = gl.getUniformLocation(program, 'uShaderEnabled');
                    gl.uniform1i(uShaderEnabled, shaderEnabled);
                }
            }
            
            // Sound meter animation
            let soundAmp = 0;
            function updateSoundMeter() {
                if (currentSource === 'browser' && !browserAudio.paused) {
                    // Simulate amplitude from browser audio
                    soundAmp = 0.3 + 0.2 * Math.sin(Date.now() * 0.01);
                    document.getElementById('soundMeter').style.width = (soundAmp * 100) + '%';
                } else if (currentSource === 'pygame') {
                    // Get amplitude from server
                    fetch('/sound/amp')
                        .then(res => res.json())
                        .then(data => {
                            soundAmp = data.amp;
                            document.getElementById('soundMeter').style.width = (soundAmp * 100) + '%';
                        });
                } else {
                    soundAmp = 0;
                    document.getElementById('soundMeter').style.width = '0%';
                }
                requestAnimationFrame(updateSoundMeter);
            }
            updateSoundMeter();
            
            function setSoundSource(source) {
                currentSource = source;
                
                // Update UI
                document.getElementById('btnNone').className = source === 'none' ? 'sound active' : 'sound';
                document.getElementById('btnPygame').className = source === 'pygame' ? 'sound active' : 'sound';
                document.getElementById('btnBrowser').className = source === 'browser' ? 'sound active' : 'sound';
                
                const indicator = document.getElementById('sourceIndicator');
                indicator.className = `indicator ${source}`;
                indicator.innerHTML = `Source: ${source.charAt(0).toUpperCase() + source.slice(1)}`;
                
                // Handle audio
                if (source === 'browser') {
                    browserAudio.play().catch(e => console.log('Audio play failed:', e));
                } else {
                    browserAudio.pause();
                    browserAudio.currentTime = 0;
                }
                
                // Tell server about source change
                fetch('/sound/source', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({source: source})
                });
            }
            
            // Mouse tracking
            canvas.addEventListener('mousemove', (e) => {
                const rect = canvas.getBoundingClientRect();
                mouse[0] = (e.clientX - rect.left) / rect.width;
                mouse[1] = 1.0 - (e.clientY - rect.top) / rect.height;
                
                const x = Math.round((e.clientX - rect.left) * 640 / rect.width);
                const y = Math.round((e.clientY - rect.top) * 480 / rect.height);
                
                fetch('/mouse', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({x, y})
                });
            });
            
            // Setup WebGL
            function createShader(type, source) {
                const shader = gl.createShader(type);
                gl.shaderSource(shader, source);
                gl.compileShader(shader);
                return shader;
            }
            
            // Compile shaders
            const vertexShader = createShader(gl.VERTEX_SHADER, 
                document.getElementById('vertex-shader').textContent);
            const fragmentShader = createShader(gl.FRAGMENT_SHADER, 
                document.getElementById('fragment-shader').textContent);
            
            // Create program
            const program = gl.createProgram();
            gl.attachShader(program, vertexShader);
            gl.attachShader(program, fragmentShader);
            gl.linkProgram(program);
            gl.useProgram(program);
            
            // Set up fullscreen quad
            const vertices = new Float32Array([
                -1, -1, 1, -1, -1, 1,
                -1, 1, 1, -1, 1, 1
            ]);
            
            const buffer = gl.createBuffer();
            gl.bindBuffer(gl.ARRAY_BUFFER, buffer);
            gl.bufferData(gl.ARRAY_BUFFER, vertices, gl.STATIC_DRAW);
            
            const position = gl.getAttribLocation(program, 'position');
            gl.enableVertexAttribArray(position);
            gl.vertexAttribPointer(position, 2, gl.FLOAT, false, 0, 0);
            
            // Get uniform locations
            const uTexture = gl.getUniformLocation(program, 'uTexture');
            const uTime = gl.getUniformLocation(program, 'uTime');
            const uMouse = gl.getUniformLocation(program, 'uMouse');
            const uResolution = gl.getUniformLocation(program, 'uResolution');
            const uShaderEnabled = gl.getUniformLocation(program, 'uShaderEnabled');
            
            gl.uniform1i(uTexture, 0);
            gl.uniform2f(uResolution, canvas.width, canvas.height);
            gl.uniform1i(uShaderEnabled, shaderEnabled);
            
            // Texture setup
            gl.bindTexture(gl.TEXTURE_2D, texture);
            gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR);
            gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.LINEAR);
            gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
            gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
            
            // Main loop
            function update() {
                fetch('/frame')
                    .then(res => res.arrayBuffer())
                    .then(buffer => {
                        // Update texture with new frame
                        gl.bindTexture(gl.TEXTURE_2D, texture);
                        gl.texImage2D(
                            gl.TEXTURE_2D, 0, gl.RGB, 640, 480, 0,
                            gl.RGB, gl.UNSIGNED_BYTE, new Uint8Array(buffer)
                        );
                        
                        // Update uniforms
                        const time = (Date.now() / 1000) - startTime;
                        gl.uniform1f(uTime, time);
                        gl.uniform2f(uMouse, mouse[0], mouse[1]);
                        
                        // Draw
                        gl.drawArrays(gl.TRIANGLES, 0, 6);
                        
                        requestAnimationFrame(update);
                    });
            }
            
            update();
        </script>
    </body>
    </html>
    ''')

@app.route('/frame')
def frame():
    """Return raw RGB bytes"""
    return Response(renderer.get_frame(), mimetype='application/octet-stream')

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

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

@app.route('/sound/amp')
def sound_amp():
    return {'amp': renderer.sound_amp}

# Serve static sound file
@app.route('/static/sound.mp3')
def serve_sound():
    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 + WebGL Shader + Sound")
    print("="*70)
    print("🌐 http://localhost:5000")
    print("\n🔮 Shader Toggle: See pure Pygame vs. effects")
    print("🔊 Sound Sources:")
    print("   • None - No sound")
    print("   • Pygame - sound.mp3 plays in Pygame (streamed)")
    print("   • Browser - sound.mp3 plays in browser")
    print("\n🎯 Orientation fixed: TOP at top, BOTTOM at bottom")
    print("   Circle color indicates sound source")
    print("   Sound meter shows activity")
    print("="*70 + "\n")
    app.run(host='0.0.0.0', port=5000, debug=False)