import pygame
import numpy as np
from flask import Flask, Response, render_template_string
from flask_sock import Sock
import time
import os
import cv2
import threading
import json

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

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}")

app = Flask(__name__)
sock = Sock(app)

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
        self.button_clicked = False
        self.sound_source = 'none'
        
    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 handle_click(self, x, y):
        button_rect = pygame.Rect(self.width-200, 120, 180, 40)
        if button_rect.collidepoint(x, y):
            self.button_clicked = not self.button_clicked
            return True
        return False
    
    def render_frame(self):
        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))
        font = pygame.font.Font(None, 24)
        
        # Draw TOP marker
        pygame.draw.rect(self.surface, (255, 100, 100), (10, 10, 100, 30))
        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 CLOCK
        current_time = time.time()
        seconds = int(current_time) % 60
        hundredths = int((current_time * 100) % 100)
        time_str = f"{seconds:02d}.{hundredths:02d}s"
        clock_text = font.render(time_str, True, (0, 255, 255))
        self.surface.blit(clock_text, (self.width-150, 40))
        
        # Draw BUTTON
        button_rect = pygame.Rect(self.width-200, 120, 180, 40)
        mouse_over = button_rect.collidepoint(self.mouse_x, self.mouse_y)
        
        if self.button_clicked:
            button_color = (0, 200, 0)
        elif mouse_over:
            button_color = (100, 100, 200)
        else:
            button_color = (80, 80, 80)
            
        pygame.draw.rect(self.surface, button_color, button_rect)
        pygame.draw.rect(self.surface, (200, 200, 200), button_rect, 2)
        
        btn_text = "✅ CLICKED!" if self.button_clicked else "🔘 CLICK ME"
        text_surf = font.render(btn_text, True, (255, 255, 255))
        text_rect = text_surf.get_rect(center=button_rect.center)
        self.surface.blit(text_surf, text_rect)
        
        # Draw circle
        circle_size = 30 + int(20 * np.sin(t * 2))
        if self.sound_source == 'pygame':
            color = (100, 255, 100)
        elif self.sound_source == 'browser':
            color = (100, 100, 255)
        else:
            color = (255, 100, 100)
            
        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))
        
        # 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_jpeg(self, quality=70):
        frame = self.render_frame()
        img = np.frombuffer(frame, dtype=np.uint8).reshape((self.height, self.width, 3))
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        _, jpeg = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, quality])
        return jpeg.tobytes()

renderer = ShaderRenderer()

# WebSocket for all interactions
@sock.route('/ws')
def websocket(ws):
    """Single WebSocket connection for all interaction"""
    while True:
        try:
            message = ws.receive()
            if not message:
                continue
                
            data = json.loads(message)
            
            if data['type'] == 'mouse':
                renderer.set_mouse(data['x'], data['y'])
                
            elif data['type'] == 'click':
                renderer.handle_click(data['x'], data['y'])
                # Broadcast button state to all clients? Optional
                
            elif data['type'] == 'sound':
                renderer.sound_source = data['source']
                
        except:
            break

# MJPEG stream (one long connection, zero API calls)
@app.route('/video.mjpeg')
def video_mjpeg():
    def generate():
        while True:
            frame = renderer.get_frame_jpeg()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
            time.sleep(1/30)
    
    return Response(
        generate(),
        mimetype='multipart/x-mixed-replace; boundary=frame'
    )

@app.route('/')
def index():
    return render_template_string('''
    <!DOCTYPE html>
    <html>
    <head>
        <title>🎮 Pygame + WebSocket</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; }
            .video-container {
                background: #000;
                border-radius: 12px;
                padding: 5px;
                margin: 20px 0;
                position: relative;
            }
            #mjpegImg {
                width: 100%;
                max-width: 640px;
                height: auto;
                border-radius: 8px;
                cursor: crosshair;
            }
            .mouse-coords {
                position: absolute;
                bottom: 10px;
                left: 10px;
                background: rgba(0,0,0,0.7);
                color: #4CAF50;
                padding: 5px 10px;
                border-radius: 20px;
                font-family: monospace;
            }
            .controls {
                background: #1a1a1a;
                border-radius: 12px;
                padding: 20px;
                margin-top: 20px;
            }
            .sound-buttons {
                display: flex;
                gap: 10px;
                justify-content: center;
                margin: 20px 0;
            }
            button {
                background: #333;
                color: white;
                border: none;
                padding: 12px 24px;
                border-radius: 8px;
                cursor: pointer;
                font-weight: bold;
            }
            button.active {
                background: #4CAF50;
                box-shadow: 0 0 20px #4CAF50;
            }
            .status {
                display: flex;
                justify-content: space-around;
                margin-top: 15px;
                padding: 10px;
                background: #222;
                border-radius: 8px;
            }
            .badge {
                padding: 5px 10px;
                border-radius: 20px;
                background: #333;
            }
            .badge.green { background: #4CAF50; }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎮 Pygame + WebSocket (Zero API)</h1>
            
            <div class="video-container">
                <img id="mjpegImg" src="/video.mjpeg" crossorigin="anonymous">
                <div id="mouseCoords" class="mouse-coords">X: 320, Y: 240</div>
            </div>
            
            <div class="controls">
                <h3>🔊 Sound Source</h3>
                <div class="sound-buttons">
                    <button id="btnNone" onclick="setSound('none')" class="active">🔇 None</button>
                    <button id="btnPygame" onclick="setSound('pygame')">🎮 Pygame</button>
                    <button id="btnBrowser" onclick="setSound('browser')">🌐 Browser</button>
                </div>
                
                <div class="status">
                    <div>Connection: <span id="wsStatus" class="badge green">🟢 Connected</span></div>
                    <div>API Calls: <span class="badge">0</span></div>
                </div>
                
                <p style="color: #666; font-size: 12px; margin-top: 15px;">
                    ⚡ Zero API polling • All interaction via WebSocket • MJPEG stream
                </p>
            </div>
        </div>
        
        <audio id="browserAudio" loop style="display:none;">
            <source src="/static/sound.mp3" type="audio/mpeg">
        </audio>
        
        <script>
            const img = document.getElementById('mjpegImg');
            const browserAudio = document.getElementById('browserAudio');
            
            // WebSocket connection (single socket for everything)
            const ws = new WebSocket((location.protocol === 'https:' ? 'wss:' : 'ws:') + '//' + window.location.host + '/ws');
            
            ws.onopen = () => document.getElementById('wsStatus').innerHTML = '🟢 Connected';
            ws.onclose = () => document.getElementById('wsStatus').innerHTML = '🔴 Disconnected';
            
            // Mouse tracking - send via WebSocket
            let mouseTimer;
            img.addEventListener('mousemove', (e) => {
                const rect = img.getBoundingClientRect();
                const x = Math.round((e.clientX - rect.left) * 640 / rect.width);
                const y = Math.round((e.clientY - rect.top) * 480 / rect.height);
                
                document.getElementById('mouseCoords').innerHTML = `X: ${x}, Y: ${y}`;
                
                // Throttle to 30fps
                if (mouseTimer) clearTimeout(mouseTimer);
                mouseTimer = setTimeout(() => {
                    ws.send(JSON.stringify({
                        type: 'mouse',
                        x: x,
                        y: y
                    }));
                }, 33);
            });
            
            // Click handling - send via WebSocket
            img.addEventListener('click', (e) => {
                const rect = img.getBoundingClientRect();
                const x = Math.round((e.clientX - rect.left) * 640 / rect.width);
                const y = Math.round((e.clientY - rect.top) * 480 / rect.height);
                
                ws.send(JSON.stringify({
                    type: 'click',
                    x: x,
                    y: y
                }));
            });
            
            // Sound handling
            function setSound(source) {
                document.getElementById('btnNone').className = source === 'none' ? 'active' : '';
                document.getElementById('btnPygame').className = source === 'pygame' ? 'active' : '';
                document.getElementById('btnBrowser').className = source === 'browser' ? 'active' : '';
                
                if (source === 'browser') {
                    browserAudio.play().catch(e => console.log('Audio error:', e));
                } else {
                    browserAudio.pause();
                    browserAudio.currentTime = 0;
                }
                
                ws.send(JSON.stringify({
                    type: 'sound',
                    source: source
                }));
            }
        </script>
    </body>
    </html>
    ''')

@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__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)