#!/usr/bin/env python3
# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Local server for development, with support for CORS headers and MIME types.

NOTE: This is NOT suitable for production serving, it is just a slightly
extended version of https://docs.python.org/3/library/http.server.html.

Usage:
  python3 local_server.py --directory {build_dir}
  (then open http://localhost:8000/ in your browser)
"""

import os
from functools import partial
from http import server


class CORSHTTPRequestHandler(server.SimpleHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        # Include MIME types for files we expect to be serving.
        # https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/MIME_types/Common_types
        self.extensions_map.update(
            {
                ".js": "application/javascript",
                ".wasm": "application/wasm",
            }
        )
        super().__init__(*args, **kwargs)

    # Inspiration for this hack: https://stackoverflow.com/a/13354482
    def end_headers(self):
        self.send_cors_headers()

        server.SimpleHTTPRequestHandler.end_headers(self)

    def send_cors_headers(self):
        # Emscripten uses SharedArrayBuffer for its multithreading, which requires
        # Cross Origin Opener Policy and Cross Origin Embedder Policy headers:
        #   * https://emscripten.org/docs/porting/pthreads.html
        #   * https://developer.chrome.com/blog/enabling-shared-array-buffer/
        self.send_header("Cross-Origin-Embedder-Policy", "require-corp")
        self.send_header("Cross-Origin-Opener-Policy", "same-origin")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--directory",
        "-d",
        default=os.getcwd(),
        help="Specify alternative directory " "[default:current directory]",
    )
    parser.add_argument(
        "port",
        action="store",
        default=8000,
        type=int,
        nargs="?",
        help="Specify alternate port [default: 8000]",
    )
    args = parser.parse_args()

    server.test(
        HandlerClass=partial(CORSHTTPRequestHandler, directory=args.directory),
        port=args.port,
    )