"""Local dev server for the demos: like `python3 -m http.server`, but never lets the
browser reuse a stale copy of pkg/*.js or *.wasm after `wasm-pack build`, and makes the
pages cross-origin isolated (COOP + COEP) so SharedArrayBuffer — and with it the
multithreaded path tracer — is available. All resources are same-origin, so COEP
blocks nothing.

Run from the repo root:  python3 scripts/serve.py [port] [--mcp mcp.config.json]
Then open http://localhost:<port>/www/

--mcp also enables the bridge for the MCP console page (www/mcp/, see scripts/mcp_bridge.py).
Its API only answers requests from this machine; the static pages stay reachable from the
LAN as before.
"""
import argparse
import functools
import http.server
import pathlib
import sys

ROOT = pathlib.Path(__file__).resolve().parent.parent
REQUEST_TIMEOUT = 60  # seconds of socket inactivity before a connection is dropped
# Pages that use the bridge's API (/api/mcp, /api/compute) and so receive its token on load
API_PAGES = ('/www/mcp/', '/www/server-compute/')


class NoCacheHandler(http.server.SimpleHTTPRequestHandler):
    extensions_map = {**http.server.SimpleHTTPRequestHandler.extensions_map,
                      '.wasm': 'application/wasm', '.js': 'text/javascript'}
    bridge = None   # set by make_server when --mcp is given
    compute = None  # … and --compute

    def end_headers(self):
        # Revalidate every request: the browser still gets 304s, but never a stale module
        self.send_header('Cache-Control', 'no-cache')
        self.send_header('Cross-Origin-Opener-Policy', 'same-origin')
        self.send_header('Cross-Origin-Embedder-Policy', 'require-corp')
        # No page may be shown in a frame (clickjacking: the MCP console can run programs)
        self.send_header('X-Frame-Options', 'DENY')
        self.send_header('Content-Security-Policy', "frame-ancestors 'none'")
        # The pages using the API, opened from this machine, receive the bridge's token
        if self.bridge and self.path.startswith(API_PAGES) and self.bridge.may_issue_cookie(self):
            self.send_header('Set-Cookie', self.bridge.cookie_header())
        super().end_headers()

    def api(self, method):
        if self.compute and self.compute.handle(self, method):
            return True
        return bool(self.bridge and self.bridge.handle(self, method))

    def do_GET(self):
        if self.path.startswith('/api/'):
            if not self.api('GET'):
                self.send_error(404, 'API not enabled (start serve.py with --mcp mcp.config.json [--compute])')
            return
        super().do_GET()

    def do_POST(self):
        if not self.api('POST'):
            self.send_error(404 if self.path.startswith('/api/') else 405)


def make_server(port, bridge=None, host='', quiet=False, compute=None):
    # A socket timeout, so a client that sends a request slowly (or never finishes it) can't
    # hold a thread forever; event streams write a keepalive well within it
    attrs = {'bridge': bridge, 'compute': compute, 'timeout': REQUEST_TIMEOUT}
    if quiet:
        attrs['log_message'] = lambda self, *args: None
    handler = type('Handler', (NoCacheHandler,), attrs)
    httpd = http.server.ThreadingHTTPServer((host, port), functools.partial(handler, directory=str(ROOT)))
    httpd.daemon_threads = True  # long-lived event streams must not block Ctrl+C
    return httpd


def main():
    ap = argparse.ArgumentParser(description='Dev server for the demos')
    ap.add_argument('port', nargs='?', type=int, default=8080)
    ap.add_argument('--mcp', metavar='CONFIG', help='enable the MCP bridge with this config file')
    ap.add_argument('--compute', action='store_true',
                    help='enable the compute API for www/server-compute/ (needs --mcp: same access rules)')
    args = ap.parse_args()
    if args.compute and not args.mcp:
        sys.exit('--compute needs --mcp mcp.config.json (it uses the bridge\'s access rules; the config may list no servers)')
    bridge = None
    if args.mcp:
        import mcp_bridge
        try:
            bridge = mcp_bridge.Bridge(args.mcp, args.port)
        except (OSError, ValueError) as e:  # missing file, bad JSON, bad config, missing password
            sys.exit(f'MCP config {args.mcp}: {e}')
    compute = None
    if args.compute:
        import compute_bridge
        compute = compute_bridge.ComputeBridge(ROOT / 'target' / 'release' / 'compute', bridge.authorize)
    with make_server(args.port, bridge, compute=compute) as httpd:
        print(f'Serving {ROOT} at http://localhost:{args.port}/www/  (Ctrl+C to stop)')
        if bridge:
            print('\n'.join(bridge.startup_lines()))
            print(f'MCP console at http://localhost:{args.port}/www/mcp/ (API answers this machine only)')
        if compute:
            print(compute.startup_line() + f'; page at http://localhost:{args.port}/www/server-compute/')
        try:
            httpd.serve_forever()
        except KeyboardInterrupt:
            pass
        finally:
            if bridge:
                bridge.shutdown()


if __name__ == '__main__':
    main()