#!/usr/bin/env python3
"""Tiny static server for v3 internal testing.

Serves v3/ over HTTP with CORS + the COOP/COEP headers ORT-Web needs for
SharedArrayBuffer / threaded WASM. Default port 8090.

Run from package/lipsync-wasm/v3/:
    python3 serve.py
Then open http://localhost:8090/demo/
"""
import atexit
import os
import socket
import subprocess
import sys
from http.server import ThreadingHTTPServer, SimpleHTTPRequestHandler

# Other tenants on this box hold many of the 80xx / 8090 ports. Pick a
# higher default and probe upwards if it's also taken. Override with PORT=.
DEFAULT_PORT = int(os.environ.get('PORT', '8001'))

# Bind to loopback only — internal test, no LAN exposure.
# Access from another host requires an SSH tunnel:
#     ssh -L 8190:localhost:8190 usr039-gpuvm01
BIND_HOST = '127.0.0.1'

# Allow .onnx + .npy + .data MIME types
MIME_OVERRIDES = {
    '.onnx': 'application/octet-stream',
    '.data': 'application/octet-stream',
    '.npy':  'application/octet-stream',
    '.json': 'application/json',
    '.js':   'application/javascript',
    '.html': 'text/html; charset=utf-8',
}


class Handler(SimpleHTTPRequestHandler):
    def guess_type(self, path):
        for ext, mime in MIME_OVERRIDES.items():
            if path.endswith(ext):
                return mime
        return super().guess_type(path)

    def end_headers(self):
        self.send_header('Cross-Origin-Opener-Policy', 'same-origin')
        self.send_header('Cross-Origin-Embedder-Policy', 'require-corp')
        self.send_header('Cross-Origin-Resource-Policy', 'cross-origin')
        # Internal-test only — wide-open CORS
        self.send_header('Access-Control-Allow-Origin', '*')
        # Always serve fresh source/HTML so JS edits land on plain refresh.
        # Big model files get a long TTL — they don't change between sessions.
        path = self.path.lower()
        if path.endswith('.onnx') or path.endswith('.data') or path.endswith('.npy'):
            self.send_header('Cache-Control', 'public, max-age=86400, immutable')
        else:
            self.send_header('Cache-Control', 'no-store, must-revalidate')
            self.send_header('Pragma', 'no-cache')
        super().end_headers()


def find_open_port(start: int, span: int = 20) -> int:
    for p in range(start, start + span):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind((BIND_HOST, p))
                return p
            except OSError:
                continue
    raise SystemExit(f"no free port in {start}..{start + span - 1}")


def spawn_ws_proxy(here: str) -> subprocess.Popen | None:
    """Start the Gemini Live WebSocket proxy as a sibling process so the
    browser can reach ws://127.0.0.1:8191/ without seeing the API key."""
    env_path = os.path.join(here, '.env')
    if not os.path.exists(env_path):
        print('[serve] no .env found — skipping Gemini Live WS proxy')
        return None
    try:
        proc = subprocess.Popen(
            [sys.executable, os.path.join(here, 'serve_ws.py')],
            cwd=here,
        )
        atexit.register(lambda: proc.terminate())
        print(f'[serve] spawned WS proxy (pid={proc.pid}) on ws://127.0.0.1:3001/')
        return proc
    except Exception as e:
        print(f'[serve] failed to spawn WS proxy: {e}')
        return None


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    os.chdir(here)
    port = find_open_port(DEFAULT_PORT)
    print(f"serving {here} on http://localhost:{port}/")
    print(f"  → http://localhost:{port}/demo/")
    print(f"   models/: {sum(1 for _ in os.scandir('models'))} files")
    spawn_ws_proxy(here)
    # ThreadingHTTPServer so the 5 parallel ONNX fetches (incl. a 117 MB
    # sidecar) don't serialize and block subsequent requests like /demo/.
    ThreadingHTTPServer((BIND_HOST, port), Handler).serve_forever()


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)