# -*- coding: utf-8 -*-
"""查看器的本地静态服务器（前端拆分后，多文件 + ES Module + Worker 都需要 http 才能跑）。

    python -m arm.viewer.serve [--port 8123] [--print-only]

要点：
  - docroot = 仓库根，`/` 直接给 webapp/index.html；/docs/*.json 这类数据同域可读。
  - Cache-Control: no-store —— 重新生成后旧标签页仍显示旧版是这里踩过的坑，
    过去只能人肉 Ctrl+F5；不缓存之后刷新即新构建。
  - 显式 MIME：Windows 上 .js 常被注册成 text/plain，那会让 <script type=module>
    静默不执行（页面全黑且没有明显报错），所以这里写死。
  - 屏蔽 .git / backups / 定位图 / out 等目录，不把整个仓库暴露到端口上。
"""
import argparse
import os
import sys
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from arm.paths import REPO

try:
    sys.stdout.reconfigure(encoding="utf-8")
except Exception:
    pass

DENY = (".git", "backups", ".bak", "定位图", "out", "scratch", "node_modules")
MIME = {
    ".js": "text/javascript; charset=utf-8",
    ".mjs": "text/javascript; charset=utf-8",
    ".json": "application/json; charset=utf-8",
    ".svg": "image/svg+xml",
    ".css": "text/css; charset=utf-8",
    ".html": "text/html; charset=utf-8",
    ".npz": "application/octet-stream",
    ".png": "image/png",
    ".wasm": "application/wasm",
}


class H(SimpleHTTPRequestHandler):
    # HTTP/1.1 + keep-alive + 大缓冲：默认 HTTP/1.0 每请求一条连接，
    # 4 个 headless Edge 同时各拉 14MB 数据脚本时会把页面加载拖到截图之后。
    protocol_version = "HTTP/1.1"
    wbufsize = 1 << 20
    # 按 (路径, mtime) 缓存字节：同一份 14MB 数据在一次批量截图里要发几十次。
    _CACHE = {}

    def __init__(self, *a, **k):
        super().__init__(*a, directory=REPO, **k)

    def _cached_bytes(self, path):
        try:
            st = os.stat(path)
        except OSError:
            return None
        key = (path, st.st_mtime_ns, st.st_size)
        body = H._CACHE.get(key)
        if body is None:
            with open(path, "rb") as f:
                body = f.read()
            if st.st_size <= 64 << 20:
                if len(H._CACHE) > 12:
                    H._CACHE.clear()
                H._CACHE[key] = body
        return body

    def send_head(self):
        rel = self.path.split("?", 1)[0].split("#", 1)[0].lstrip("/").replace("\\", "/")
        top = rel.split("/")[0].lower()
        if top in [d.lower() for d in DENY] or ".." in rel:
            self.send_error(404, "denied")
            return None
        fs = self.translate_path(self.path)
        if os.path.isfile(fs):
            body = self._cached_bytes(fs)
            if body is not None:
                self.send_response(200)
                self.send_header("Content-Type", self.guess_type(fs))
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
                return None
        return super().send_head()

    def guess_type(self, path):
        ext = os.path.splitext(path)[1].lower()
        return MIME.get(ext) or super().guess_type(path)


    def do_GET(self):
        """`/` 必须 302 到 /webapp/index.html，不能在原地把内容发出去：
        那样文档基址是 `/`，页面里所有相对路径（styles/… js/… vendor/…）都会
        解析到 /styles/base.css 之类而 404 —— 实测整页样式与脚本全丢。"""
        if self.path.split("?", 1)[0].rstrip("/") in ("", "/"):
            q = self.path.split("?", 1)[1] if "?" in self.path else ""
            self.send_response(302)
            self.send_header("Location", "/webapp/index.html" + ("?" + q if q else ""))
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        super().do_GET()

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

    def log_message(self, fmt, *args):
        sys.stderr.write("%s %s\n" % (self.address_string(), fmt % args))


def serve(port=8123):
    srv = None
    for p in range(port, port + 20):
        try:
            srv = ThreadingHTTPServer(("127.0.0.1", p), H)
            port = p
            break
        except OSError:
            continue
    if srv is None:
        raise SystemExit("端口 %d..%d 都被占用" % (port, port + 19))
    url = "http://127.0.0.1:%d/" % port
    print("docroot %s" % REPO, flush=True)
    print("查看器: %s  (Ctrl+C 停止；已禁用缓存，改完刷新即新构建)" % url, flush=True)
    srv.serve_forever()


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="查看器本地静态服务器")
    ap.add_argument("--port", type=int, default=8123)
    serve(ap.parse_args().port)