#!/usr/bin/env python3
"""Serve the offline mirror of https://gradio.app/docs.

    python3 serve.py                      # http://127.0.0.1:8091/docs
    python3 serve.py --port 9000
    python3 serve.py --host 0.0.0.0       # reachable from other machines

Files under site/ are kept exactly as gradio.app served them.  Every URL that
would leave the machine is rewritten here, while the response is streamed:

  * fonts.googleapis.com / fonts.gstatic.com  -> /_ext/<host>/...
  * the Tailwind browser build on jsDelivr    -> /_ext/cdn.jsdelivr.net/...
  * images on raw.githubusercontent.com etc.  -> /_ext/<host>/...  (only inside
    src=/srcset= attributes, so URLs quoted in code samples stay authentic)
  * gradio.s3-us-west-2.amazonaws.com/<v>/gradio.js -> /_stub/gradio-app.js
  * www.googletagmanager.com/gtag/js                -> /_stub/noop.js
"""

import argparse
import gzip
import hashlib
import io
import json
import os
import posixpath
import re
import sys
import urllib.parse
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

ROOT = os.path.dirname(os.path.abspath(__file__))
SITE = os.path.join(ROOT, "site")

MIME = {
    ".html": "text/html; charset=utf-8",
    ".json": "application/json; charset=utf-8",
    ".js": "text/javascript; charset=utf-8",
    ".mjs": "text/javascript; charset=utf-8",
    ".css": "text/css; charset=utf-8",
    ".svg": "image/svg+xml",
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".gif": "image/gif",
    ".webp": "image/webp",
    ".ico": "image/x-icon",
    ".woff": "font/woff",
    ".woff2": "font/woff2",
    ".ttf": "font/ttf",
    ".eot": "application/vnd.ms-fontobject",
    ".mp4": "video/mp4",
    ".webm": "video/webm",
    ".mov": "video/quicktime",
    ".wasm": "application/wasm",
    ".map": "application/json",
    ".txt": "text/plain; charset=utf-8",
    ".whl": "application/octet-stream",
    ".zip": "application/zip",
}
TEXT_EXT = {".html", ".js", ".mjs", ".css", ".json", ".svg", ".map", ".txt"}
TEXT_MIME = ("text/", "application/json", "application/javascript", "image/svg")

with open(os.path.join(SITE, "_ext", "manifest.json"), "rb") as _f:
    MANIFEST = json.load(_f)

# Files whose name carries no usable extension (the Google Fonts "css2?..."
# reply, for one) keep the content type gradio.app / the CDN reported.
try:
    with open(os.path.join(SITE, "_ext", "types.json"), "rb") as _f:
        TYPES = json.load(_f)
except FileNotFoundError:
    TYPES = {}

INFRA_HOSTS = MANIFEST["infra_hosts"]
MEDIA_HOSTS = MANIFEST["media_hosts"]
EXACT = dict(MANIFEST["exact"])
EXACT["https://gradio.s3-us-west-2.amazonaws.com"] = "/_stub/s3"
EXACT["https://www.googletagmanager.com/gtag/js"] = "/_stub/noop.js"
ROUTES = set(MANIFEST["routes"])

# src="https://host/..." / srcset="..." -> local copy under /_ext
MEDIA_ATTR_RE = re.compile(
    r'(?P<attr>\b(?:src|srcset|poster)=")(?P<url>https://(?:%s)/[^"]+)'
    % "|".join(re.escape(h) for h in MEDIA_HOSTS).replace("/", r"\/")
) if MEDIA_HOSTS else None

# iframes of live Hugging Face Spaces (the theming page uses ten of them)
SPACE_IFRAME_RE = re.compile(
    r'(?P<attr>\bsrc=")(?P<url>https://[A-Za-z0-9._-]+\.hf\.space[^"]*)')


def ext_local(url):
    """https://host/p?q -> /_ext/host/p[@hash]   (mirror.py's naming)"""
    u = urllib.parse.urlsplit(url.replace("&amp;", "&"))
    p = u.path.lstrip("/") or "index"
    if u.query:
        p += "@" + hashlib.sha1(u.query.encode()).hexdigest()[:12]
    return "/_ext/%s/%s" % (u.netloc, p)


def rewrite(data, ext):
    """Point every external URL in a text response at a local file."""
    try:
        text = data.decode("utf-8")
    except UnicodeDecodeError:
        return data

    for url, local in EXACT.items():
        if url == "https://gradio.s3-us-west-2.amazonaws.com":
            # .../<version>/gradio.js -> the local web-component stub
            text = re.sub(r'https://gradio\.s3-us-west-2\.amazonaws\.com/[^"\']*?/gradio\.js',
                          "/_stub/gradio-app.js", text)
        else:
            text = text.replace(url, local)

    for host in INFRA_HOSTS:
        text = text.replace("https://" + host, "/_ext/" + host)

    # component JS also embeds <img src="https://raw.githubusercontent.com/...">
    if MEDIA_ATTR_RE is not None and ext in (".html", ".svg", ".json", ".js", ".mjs", ".css"):
        text = MEDIA_ATTR_RE.sub(
            lambda m: m.group("attr") + ext_local(m.group("url")), text)

    # The theming page frames live apps on Hugging Face Spaces; point those
    # iframes at a local notice instead of an unreachable host.
    if ext in (".html", ".js", ".mjs"):
        text = SPACE_IFRAME_RE.sub(
            lambda m: m.group("attr") + "/_stub/space?u="
            + urllib.parse.quote(m.group("url").replace("&amp;", "&"), safe=""),
            text)

    return text.encode("utf-8")


OFFLINE_PAGE = """<!doctype html><html lang="en"><head><meta charset="utf-8">
<title>Not in this mirror</title><style>
:root{color-scheme:light dark;--fg:#1f2937;--bg:#fff;--mut:#6b7280;--acc:#f97316}
@media (prefers-color-scheme:dark){:root{--fg:#e5e7eb;--bg:#111827;--mut:#9ca3af}}
body{margin:0;min-height:100vh;display:flex;align-items:center;justify-content:center;
background:var(--bg);color:var(--fg);font:16px/1.6 -apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif}
main{max-width:34rem;padding:2rem}h1{font-size:1.4rem;margin:0 0 .6rem}
code{background:rgba(127,127,127,.15);padding:.1rem .35rem;border-radius:.25rem;font-size:.9em}
p{color:var(--mut)}a{color:var(--acc)}</style></head><body><main>
<h1>Outside this offline mirror</h1>
<p><code>%(path)s</code> is part of gradio.app but not of the
<strong>/docs</strong> section that was mirrored, so it is not available offline.</p>
<p><a href="/docs">&#8592; Back to the API documentation</a></p>
</main></body></html>"""


_search_cache = []


def search_index():
    """The payload the docs search worker loads from /search-api.

    gradio.app returns 232 entries whose slugs are version-prefixed
    ("/main/docs/gradio/interface") and include the guides.  Point the docs
    entries at this mirror and drop anything it does not hold, so every hit
    the search box offers actually opens.
    """
    if _search_cache:
        return _search_cache[0]
    src = os.path.join(SITE, "search-api.json")
    if not os.path.isfile(src):
        return b"[]"
    with open(src, "rb") as f:
        entries = json.load(f)
    kept = []
    for e in entries:
        slug = re.sub(r"^/main/docs/", "/docs/", e.get("slug", "")).rstrip("/")
        if slug in ROUTES:
            kept.append(dict(e, slug=slug))
    body = json.dumps(kept).encode()
    _search_cache.append(body)
    return body


SPACE_PAGE = """<!doctype html><html lang="en"><head><meta charset="utf-8">
<title>Live app unavailable offline</title><style>
:root{color-scheme:light dark}
body{margin:0;height:100vh;display:flex;align-items:center;justify-content:center;
font:14px/1.5 -apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif;
color:#6b7280;background:transparent;text-align:center}
div{border:1px dashed currentColor;border-radius:.5rem;padding:1.25rem 1.5rem;opacity:.85}
code{font-size:.9em}</style></head><body><div>
<strong>Live app not available offline</strong><br>
This frame shows <code>%(host)s</code>, a Hugging Face Space.
</div></body></html>"""


class Handler(BaseHTTPRequestHandler):
    server_version = "gradio-doc-offline/1.0"
    protocol_version = "HTTP/1.1"

    # ---------------------------------------------------------------- helpers
    def log_message(self, fmt, *args):
        if self.server.verbose:
            sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args))

    def resolve(self, path, query):
        """URL path -> file on disk, or None."""
        p = posixpath.normpath(urllib.parse.unquote(path))
        if p.startswith("//"):
            p = p[1:]
        if ".." in p.split("/"):
            return None
        rel = p.lstrip("/")

        # query strings are folded into the filename for mirrored externals
        if query:
            cand = os.path.join(
                SITE, rel + "@" + hashlib.sha1(query.encode()).hexdigest()[:12])
            if os.path.isfile(cand):
                return cand

        cand = os.path.join(SITE, rel)
        if os.path.isfile(cand):
            return cand
        if os.path.isdir(cand):
            idx = os.path.join(cand, "index.html")
            if os.path.isfile(idx):
                return idx
        # extensionless route saved as <route>/index.html
        if not os.path.splitext(rel)[1]:
            idx = os.path.join(SITE, rel, "index.html")
            if os.path.isfile(idx):
                return idx
        return None

    def send_payload(self, body, ctype, status=HTTPStatus.OK, cache=None):
        headers = [("Content-Type", ctype)]
        if cache:
            headers.append(("Cache-Control", cache))
        if (ctype.startswith(("text/", "application/json", "image/svg"))
                and len(body) > 1400
                and "gzip" in (self.headers.get("Accept-Encoding") or "")):
            buf = io.BytesIO()
            with gzip.GzipFile(fileobj=buf, mode="wb", compresslevel=6) as gz:
                gz.write(body)
            body = buf.getvalue()
            headers.append(("Content-Encoding", "gzip"))
            headers.append(("Vary", "Accept-Encoding"))
        self.send_response(status)
        for k, v in headers:
            self.send_header(k, v)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        if self.command != "HEAD":
            self.wfile.write(body)

    def redirect(self, to):
        self.send_response(HTTPStatus.FOUND)
        self.send_header("Location", to)
        self.send_header("Content-Length", "0")
        self.end_headers()

    # ------------------------------------------------------------------ verbs
    def do_HEAD(self):
        self.do_GET()

    def do_GET(self):
        path, _, query = self.path.partition("?")

        if path in ("/", "/index.html"):
            return self.redirect("/docs")

        # gradio.app 307s "/docs/gradio/interface/" to the slash-less URL, and
        # the sidebar links are relative ("audio"), so they only resolve from
        # the canonical form.  Keep that behaviour.
        if len(path) > 1 and path.endswith("/"):
            canonical = path.rstrip("/") or "/"
            return self.redirect(canonical + (("?" + query) if query else ""))

        # locally generated helpers
        if path == "/_stub/gradio-app.js":
            return self.send_payload(GRADIO_STUB.encode(),
                                     MIME[".js"], cache="no-cache")
        if path == "/_stub/noop.js":
            return self.send_payload(b"/* analytics disabled offline */\n",
                                     MIME[".js"], cache="no-cache")
        if path == "/_stub/space":
            space = urllib.parse.parse_qs(query).get("u", [""])[0]
            host = urllib.parse.urlsplit(space).netloc or "a Hugging Face Space"
            body = (SPACE_PAGE % {"host": host.replace("<", "&lt;")}).encode()
            return self.send_payload(body, MIME[".html"], cache="no-cache")
        if path == "/search-api":
            return self.send_payload(search_index(), MIME[".json"],
                                     cache="no-cache")

        target = self.resolve(path, query)
        if target is None:
            # SvelteKit data request for a route we do not have
            if path.endswith("__data.json"):
                return self.send_payload(b'{"type":"error"}', MIME[".json"],
                                         HTTPStatus.NOT_FOUND)
            body = (OFFLINE_PAGE % {"path": path.replace("<", "&lt;")}).encode()
            return self.send_payload(body, MIME[".html"], HTTPStatus.NOT_FOUND)

        ext = os.path.splitext(target)[1].lower()
        rel = os.path.relpath(target, SITE)
        ctype = MIME.get(ext) or TYPES.get(rel) or "application/octet-stream"
        if ctype.startswith(("text/", "application/json")) and "charset" not in ctype:
            ctype += "; charset=utf-8"
        with open(target, "rb") as f:
            body = f.read()
        if ext in TEXT_EXT or ctype.startswith(TEXT_MIME):
            body = rewrite(body, ext or os.path.splitext(rel.split("@")[0])[1].lower()
                           or (".css" if ctype.startswith("text/css") else ""))

        cache = ("public, max-age=31536000, immutable"
                 if path.startswith("/_app/immutable/") else "no-cache")
        self.send_payload(body, ctype, cache=cache)


GRADIO_STUB = r"""/* Offline stand-in for gradio.js (the <gradio-app> web component).
   The real component streams a live demo from a Hugging Face Space, which is
   unreachable without internet access, so it is replaced by a notice.  The
   source code of every demo is still shown above it on the page. */
(function () {
  var CSS = "display:block;border:1px dashed currentColor;border-radius:.5rem;" +
    "opacity:.7;padding:1.25rem;margin:.5rem 0;font:14px/1.5 -apple-system," +
    "BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif;text-align:center";
  function render(el) {
    var space = el.getAttribute("space") || el.getAttribute("src") || "";
    el.setAttribute("style", CSS);
    el.innerHTML =
      '<strong>Live demo not available offline</strong><br>' +
      (space ? 'Hosted at <code>' + space.replace(/[<>&"]/g, "") + '</code> ' +
               'on Hugging Face Spaces.<br>' : '') +
      'Run the demo code shown above locally instead.';
  }
  function define(tag) {
    if (window.customElements && !customElements.get(tag)) {
      customElements.define(tag, class extends HTMLElement {
        connectedCallback() { render(this); }
      });
    }
  }
  define("gradio-app");
  define("gradio-lite");
})();
"""


def main():
    ap = argparse.ArgumentParser(description="Serve the offline gradio.app/docs mirror")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8091)
    ap.add_argument("-v", "--verbose", action="store_true", help="log every request")
    args = ap.parse_args()

    if not os.path.isdir(os.path.join(SITE, "docs")):
        sys.exit("site/docs is missing - run tools/mirror.py first")

    httpd = ThreadingHTTPServer((args.host, args.port), Handler)
    httpd.verbose = args.verbose
    shown = args.host if args.host != "0.0.0.0" else "<this-host>"
    print("Gradio documentation (offline)  ->  http://%s:%d/docs" % (shown, args.port))
    print("Ctrl-C to stop")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\nstopped")


if __name__ == "__main__":
    main()