#!/usr/bin/env python3
"""
OpenCV WASM Test Server

Serves static files + API endpoints:
  GET  /api/version  - Run opencv_version.js, return version string
  POST /api/awb      - Run Gray World AWB via awb_static.wasm
"""
import http.server
import json
import os
import subprocess
import sys
import urllib.parse
import base64

PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 8910
DIR = os.path.dirname(os.path.abspath(__file__))


class Handler(http.server.SimpleHTTPRequestHandler):

    def do_GET(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == '/api/version':
            self.handle_version()
        else:
            super().do_GET()

    def do_POST(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == '/api/awb':
            self.handle_awb()
        else:
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b'{"success": false, "output": "Not found"}')

    def handle_version(self):
        """Run opencv_version.wasm via Node.js (using locateFile wrapper)"""
        self.send_response(200)
        self.send_header('Content-Type', 'application/json')
        self.end_headers()

        js_path = os.path.join(DIR, 'opencv_version.js')
        lib_dir = os.path.join(DIR, 'lib')

        # Node.js wrapper with locateFile config
        wrapper_script = f"""
const mod = require({json.dumps(js_path)});
const path = require('path');
const fs = require('fs');
mod.locateFile = function(fp, prefix) {{
    const lib = {json.dumps(lib_dir)};
    if (fp.includes('.so')) {{
        const exact = path.join(lib, fp);
        if (fs.existsSync(exact)) return exact;
        const base = fp.replace(/\\.\\d+\\.\\d+\\.\\d+$/, '').replace(/\\.\\d+$/, '');
        const noVer = path.join(lib, base);
        if (fs.existsSync(noVer)) return noVer;
    }}
    return path.join(prefix, fp);
}};
"""
        try:
            result = subprocess.run(
                ['node', '-e', wrapper_script],
                capture_output=True,
                text=True,
                timeout=30,
            )
            output = result.stdout.strip() or result.stderr.strip()
            response = {
                'success': result.returncode == 0,
                'output': output,
                'returncode': result.returncode
            }
        except subprocess.TimeoutExpired:
            response = {'success': False, 'output': 'Timeout', 'returncode': -1}
        except Exception as e:
            response = {'success': False, 'output': str(e), 'returncode': -1}

        self.wfile.write(json.dumps(response, ensure_ascii=False).encode('utf-8'))

    def handle_awb(self):
        """Run Gray World AWB via awb_static.wasm

        POST body: JSON with {"image": "<base64_of_input_png>"}
        Response: JSON with {"success": true, "image": "<base64_of_output_png>",
                             "metadata": "<stderr_output>"}
        """
        content_length = int(self.headers.get('Content-Length', 0))
        if content_length == 0:
            self.send_error_response('Empty request body')
            return

        body = self.rfile.read(content_length)
        try:
            data = json.loads(body)
        except json.JSONDecodeError:
            self.send_error_response('Invalid JSON')
            return

        input_b64 = data.get('image', '')
        if not input_b64:
            self.send_error_response('Missing image field')
            return

        # Validate base64 image
        try:
            raw = base64.b64decode(input_b64)
            if raw[:4] != b'\x89PNG':
                self.send_error_response('Only PNG images are supported')
                return
        except Exception:
            self.send_error_response('Invalid base64 image')
            return

        js_path = os.path.join(DIR, 'awb_static.js')

        # Node.js wrapper - runs AWB with locateFile (not needed for static build)
        wrapper_script = f"""
const mod = require({json.dumps(js_path)});
mod.noInitialRun = true;
mod.preRun = function() {{
    const buf = Buffer.from(process.argv[1], 'base64');
    mod.FS.mkdirTree('/data');
    mod.FS.writeFile('/data/input.png', new Uint8Array(buf));
}};
mod.onRuntimeInitialized = function() {{
    try {{
        mod.callMain([]);
        const outputBytes = mod.FS.readFile('/data/output.png');
        process.stdout.write(Buffer.from(outputBytes).toString('base64'));
    }} catch (e) {{
        process.stderr.write('ERR:' + (e.message || e));
        process.exit(1);
    }}
}};
"""

        try:
            result = subprocess.run(
                ['node', '-e', wrapper_script, input_b64],
                capture_output=True,
                text=False,  # binary mode - preserve base64 encoding
                timeout=60,
            )
            stdout_b64 = result.stdout.decode('utf-8', errors='replace').strip()
            stderr_txt = result.stderr.decode('utf-8', errors='replace').strip()

            response = {
                'success': result.returncode == 0 and len(stdout_b64) > 50,
                'image': stdout_b64,
                'metadata': stderr_txt,
                'returncode': result.returncode,
            }
        except subprocess.TimeoutExpired:
            response = {'success': False, 'error': 'Timeout processing image'}
        except Exception as e:
            response = {'success': False, 'error': str(e)}

        self.send_response(200)
        self.send_header('Content-Type', 'application/json')
        self.end_headers()
        self.wfile.write(json.dumps(response, ensure_ascii=False).encode('utf-8'))

    def send_error_response(self, msg):
        self.send_response(400)
        self.send_header('Content-Type', 'application/json')
        self.end_headers()
        self.wfile.write(json.dumps(
            {'success': False, 'error': msg}, ensure_ascii=False
        ).encode('utf-8'))

    def log_message(self, fmt, *args):
        sys.stderr.write(f'[server] {args[0]} {args[1]} {args[2]}\n')


if __name__ == '__main__':
    http.server.SimpleHTTPRequestHandler.extensions_map.update({
        '.wasm': 'application/wasm',
        '.so': 'application/octet-stream',
    })

    server = http.server.HTTPServer(('0.0.0.0', PORT), Handler)
    print(f'🚀 OpenCV WASM Test Server running at http://0.0.0.0:{PORT}')
    print(f'   📁 Static files: {DIR}')
    print(f'   📡 API: /api/version | POST /api/awb')
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        server.shutdown()
        print('\n👋 Server stopped')