"""Static server for the dev preview, with configurable COOP/COEP headers.

Usage: serve.py DIR [PORT]   (serves on 127.0.0.1)
Env:   ISOLATE=0           omit the headers (for A/B tests of cross-origin isolation)
       ISOLATE_PREFIX=/x   only send them under this URL prefix (default: everywhere)
       ONNX_BYTES_PER_SECOND  throttle model downloads (default: 0, unlimited)
"""
import functools
import os
import http.server
import mimetypes
import sys
import time

mimetypes.add_type("text/javascript", ".mjs")
mimetypes.add_type("application/wasm", ".wasm")
mimetypes.add_type("application/manifest+json", ".webmanifest")


class Handler(http.server.SimpleHTTPRequestHandler):
    def copyfile(self, source, outputfile):
        rate = int(os.environ.get("ONNX_BYTES_PER_SECOND", "0"))
        if rate <= 0 or not self.path.split("?")[0].endswith(".onnx"):
            return super().copyfile(source, outputfile)
        try:
            while chunk := source.read(min(65536, rate)):
                outputfile.write(chunk)
                outputfile.flush()
                time.sleep(len(chunk) / rate)
        except (BrokenPipeError, ConnectionResetError):
            pass  # a failed/cancelled download is expected in recovery checks

    def end_headers(self):
        if os.environ.get("ISOLATE", "1") != "0" and self.path.startswith(os.environ.get("ISOLATE_PREFIX", "/")):
            self.send_header("Cross-Origin-Opener-Policy", "same-origin")
            self.send_header("Cross-Origin-Embedder-Policy", "require-corp")
        super().end_headers()

    def log_message(self, *args):
        pass


if __name__ == "__main__":
    port = int(sys.argv[2]) if len(sys.argv) > 2 else 8765
    handler = functools.partial(Handler, directory=sys.argv[1])
    http.server.ThreadingHTTPServer(("127.0.0.1", port), handler).serve_forever()