#!/usr/bin/env python3
"""
Local server for the wllama voice-to-voice Space.

Mirrors the Hugging Face Space headers (COOP/COEP) so cross-origin
isolation, multithreading and the OPFS model cache work locally, and
serves .wasm with the correct MIME type.

Usage:
    python3 serve.py [port]     # default port 8000
"""
import os
import sys
import http.server
import socketserver

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

MIME = {
    ".html": "text/html; charset=utf-8",
    ".js": "text/javascript; charset=utf-8",
    ".mjs": "text/javascript; charset=utf-8",
    ".css": "text/css; charset=utf-8",
    ".wasm": "application/wasm",
    ".json": "application/json; charset=utf-8",
    ".onnx": "application/octet-stream",
    ".data": "application/octet-stream",
    ".png": "image/png",
    ".ico": "image/x-icon",
    ".svg": "image/svg+xml",
}


class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=ROOT, **kwargs)

    def end_headers(self):
        # Cross-origin isolation: required for SharedArrayBuffer / multithreading / OPFS
        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")
        self.send_header("Cache-Control", "no-store")
        super().end_headers()

    def guess_type(self, path):
        _, ext = os.path.splitext(path)
        return MIME.get(ext, "application/octet-stream")


class ThreadingServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
    daemon_threads = True


if __name__ == "__main__":
    print(f"Serving {ROOT} at http://localhost:{PORT}/  (Ctrl+C to stop)")
    with ThreadingServer(("", PORT), Handler) as httpd:
        httpd.serve_forever()