#!/usr/bin/env python3
"""Serve the offline Semgrep docs mirror in ./site.

Usage: python3 serve.py [--host 127.0.0.1] [--port 8080]

- Resolves extension-less page URLs (/writing-rules/overview -> overview.html),
  follows redirects recorded by the mirror, and falls back to 404.html.
- The mirror keeps files byte-for-byte as downloaded. When serving HTML/JS/CSS/JSON,
  CDN URLs (https://mintcdn.com/...) are rewritten to absolute URLs on this server
  (http://<Host header>/_ext/mintcdn.com/...). They must stay absolute: the site's
  JavaScript calls new URL() on them and crashes on relative paths.

Stdlib only; runs on Python 3.9+.
"""
import argparse
import http.server
import io
import json
import os
import posixpath
import re
import threading
import urllib.parse

SITE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "site")
CDN_HOSTS = (
    "mintcdn.com",
    "d3gk2c5xim1je2.cloudfront.net",
    "d4tuoctqmanu0.cloudfront.net",
    "mintlify.b-cdn.net",
)
CDN_RE = re.compile(rb"https:(\\?/)\\?/(" + b"|".join(re.escape(h.encode()) for h in CDN_HOSTS) + rb")")
CDN_HTML_RE = re.compile(rb"https://(" + b"|".join(re.escape(h.encode()) for h in CDN_HOSTS) + rb")")
REWRITE_EXT = (".html", ".js", ".mjs", ".css", ".json")
SAFE_HOST_RE = re.compile(r"^[A-Za-z0-9.\-]+(:\d{1,5})?$|^\[[0-9A-Fa-f:.]+\](:\d{1,5})?$")

try:
    with open(os.path.join(SITE, "_redirects.json")) as f:
        REDIRECTS = json.load(f)
except OSError:
    REDIRECTS = {}

_cache = {}
_cache_lock = threading.Lock()


def rewritten(fp, origin):
    """File contents with CDN URLs pointed at origin; cached per (file, mtime, origin)."""
    key = (fp, os.stat(fp).st_mtime_ns, origin)
    with _cache_lock:
        if key in _cache:
            return _cache[key]
    with open(fp, "rb") as f:
        data = f.read()
    if fp.endswith(".html"):
        # HTML embeds the React Server Components stream, whose text rows are
        # length-prefixed: replacements must keep the byte length unchanged.
        # "https://" and "/./_ext/" are both 8 bytes; browsers normalize "/./".
        data = CDN_HTML_RE.sub(lambda m: b"/./_ext/" + m.group(1), data)
    else:
        o = origin.encode()
        data = CDN_RE.sub(lambda m: b"http:%s%s%s%s_ext%s%s" % (
            m.group(1), m.group(1), o.replace(b"/", m.group(1)), m.group(1), m.group(1), m.group(2)), data)
    with _cache_lock:
        if len(_cache) > 400:
            _cache.clear()
        _cache[key] = data
    return data


class Handler(http.server.SimpleHTTPRequestHandler):
    extensions_map = dict(
        http.server.SimpleHTTPRequestHandler.extensions_map,
        **{".js": "text/javascript", ".mjs": "text/javascript", ".css": "text/css",
           ".svg": "image/svg+xml", ".woff2": "font/woff2", ".woff": "font/woff",
           ".json": "application/json", ".webp": "image/webp", ".html": "text/html; charset=utf-8"}
    )

    def __init__(self, *a, **kw):
        super().__init__(*a, directory=SITE, **kw)

    def origin(self):
        host = self.headers.get("Host", "")
        if not SAFE_HOST_RE.match(host):
            h, p = self.server.server_address[:2]
            host = "%s:%d" % ("127.0.0.1" if h in ("0.0.0.0", "") else h, p)
        return host

    def resolve(self):
        path = urllib.parse.unquote(urllib.parse.urlsplit(self.path).path)
        path = posixpath.normpath(path)
        if ".." in path.split("/"):
            return None, None
        key = path.rstrip("/") or "/"
        rel = key.lstrip("/")
        for cand in (rel, rel + ".html", posixpath.join(rel, "index.html")):
            fp = os.path.join(SITE, cand)
            if os.path.isfile(fp):
                return "file", fp
        if key in REDIRECTS:
            return "redirect", REDIRECTS[key]
        return None, None

    def send_file(self, code, fp):
        if fp.endswith(REWRITE_EXT):
            body = rewritten(fp, self.origin())
        else:
            with open(fp, "rb") as f:
                body = f.read()
        self.send_response(code)
        self.send_header("Content-Type", self.guess_type(fp))
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        return io.BytesIO(body)

    def send_head(self):
        kind, target = self.resolve()
        if kind == "file":
            return self.send_file(200, target)
        if kind == "redirect":
            self.send_response(301)
            self.send_header("Location", target)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return None
        fp = os.path.join(SITE, "404.html")
        if os.path.isfile(fp):
            return self.send_file(404, fp)
        self.send_error(404)
        return None

    def end_headers(self):
        self.send_header("Cache-Control", "no-cache")
        super().end_headers()


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--host", default="127.0.0.1", help="bind address (0.0.0.0 for all interfaces)")
    ap.add_argument("--port", type=int, default=8080)
    args = ap.parse_args()
    if not os.path.isfile(os.path.join(SITE, "index.html")):
        raise SystemExit("mirror not found: %s/index.html" % SITE)
    srv = http.server.ThreadingHTTPServer((args.host, args.port), Handler)
    print("Semgrep docs offline: http://%s:%d/" % (args.host, args.port), flush=True)
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()