"""Stdlib regressions for build, publish and preview. Run: python3 test/site_tools.py."""
import contextlib
import hashlib
import json
import os
from pathlib import Path
import re
import signal
import socket
import subprocess
import tempfile
import time
import unittest
from urllib.error import HTTPError, URLError
from urllib.request import ProxyHandler, build_opener

HTTP = build_opener(ProxyHandler({}))
REPO = Path(__file__).resolve().parents[1]
SCRIPT_ENV = """RELEASE BASE MODELS_URL RUNTIME_URL SOURCE CF_ANALYTICS_TOKEN MODEL_ID SIZE
STATIC_ROOT FILE_MODE DIR_MODE MODEL RUNTIME OLDER PORT NEWS FEATURES TRAINED
ISOLATE ISOLATE_PREFIX ONNX_BYTES_PER_SECOND BASH_ENV ENV""".split()


class SiteTools(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # site/build.sh compiles the shared rules to WebAssembly; a cold Rust
        # build is slow, so build them once here and keep later runs cached.
        subprocess.run([str(REPO / "core/wasm/build.sh")], cwd=REPO, check=True, timeout=900)

    def setUp(self):
        temporary = tempfile.TemporaryDirectory()
        self.addCleanup(temporary.cleanup)
        self.root = Path(temporary.name)
        self.env = os.environ.copy()
        for key in SCRIPT_ENV:
            self.env.pop(key, None)

    def model(self, name, content=b"fixture model\n"):
        path = self.root / name
        (path / "onnx").mkdir(parents=True)
        (path / "config.json").write_text('{"model_type": "fixture"}\n')
        (path / "onnx/model_q4.onnx").write_bytes(content)
        (path / "SHA256SUMS").write_text("".join(
            f"{hashlib.sha256((path / file).read_bytes()).hexdigest()}  {file}\n"
            for file in ("config.json", "onnx/model_q4.onnx")))
        return path

    def runtime(self, name, label):
        path = self.root / name
        path.mkdir(parents=True)
        (path / "transformers.min.js").write_text(f"export const fixture = '{label}';\n")
        (path / "ort.wasm").write_bytes(label.encode())
        return path

    def tree(self, path):
        result = {}
        for entry in path.rglob("*"):
            name = entry.relative_to(path).as_posix()
            result[name] = (("link", os.readlink(entry)) if entry.is_symlink() else
                            ("dir",) if entry.is_dir() else ("file", entry.read_bytes()))
        return result

    def run_script(self, name, *args, success=True, **env):
        # A timed-out script is stopped with all of its children (cargo,
        # esbuild), so it cannot keep racing the next test's build.
        process = subprocess.Popen([str(REPO / name), *map(str, args)], cwd=REPO,
                                   env={**self.env, **env}, stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE, text=True, start_new_session=True)
        try:
            stdout, stderr = process.communicate(timeout=120)
        except subprocess.TimeoutExpired:
            with contextlib.suppress(ProcessLookupError):
                os.killpg(process.pid, signal.SIGKILL)
            process.communicate()
            raise
        result = subprocess.CompletedProcess(process.args, process.returncode, stdout, stderr)
        message = result.stdout + result.stderr
        if success:
            self.assertEqual(result.returncode, 0, message)
        else:
            self.assertNotEqual(result.returncode, 0, message)
        return result

    def config(self, page):
        match = re.search(r'id="pet-config">(.*?)</script>', page, re.S)
        self.assertIsNotNone(match)
        return json.loads(match.group(1))

    def assert_manifest(self, path):
        for line in (path / "SHA256SUMS").read_text().splitlines():
            digest, name = line.split("  ", 1)
            self.assertEqual(hashlib.sha256((path / name).read_bytes()).hexdigest(), digest)

    @staticmethod
    def fetch(url):
        with HTTP.open(url, timeout=1) as response:
            return response.read(), response.headers

    @contextlib.contextmanager
    def preview(self, output, model, runtime):
        with socket.socket() as listener:
            listener.bind(("127.0.0.1", 0))
            port = listener.getsockname()[1]
        url = f"http://127.0.0.1:{port}"
        with tempfile.TemporaryFile(mode="w+") as log:
            process = subprocess.Popen([str(REPO / "test/dev-site.sh"), str(output)], cwd=REPO,
                                       env={**self.env, "MODEL": str(model), "RUNTIME": str(runtime),
                                            "BASE": "pet", "MODEL_ID": "demo", "PORT": str(port)},
                                       stdout=log, stderr=subprocess.STDOUT, start_new_session=True)
            try:
                deadline = time.monotonic() + 60
                while time.monotonic() < deadline:
                    if process.poll() is not None:
                        log.seek(0)
                        self.fail(f"preview exited with {process.returncode}: {log.read()}")
                    try:
                        self.fetch(url + "/pet/")
                        break
                    except (URLError, TimeoutError):
                        time.sleep(0.05)
                else:
                    log.seek(0)
                    self.fail(f"preview did not start: {log.read()}")
                yield url
            finally:
                if process.poll() is None:
                    with contextlib.suppress(ProcessLookupError):
                        os.killpg(process.pid, signal.SIGTERM)
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    with contextlib.suppress(ProcessLookupError):
                        os.killpg(process.pid, signal.SIGKILL)
                    process.wait(timeout=5)

    def test_failed_build_preserves_rendered_site_and_published_directories(self):
        output = self.root / "site"
        self.model("site/models/demo")
        self.runtime("site/runtime/fixture-runtime", "preserved")
        published = {name: self.tree(output / name) for name in ("models", "runtime")}
        release = {"modelId": "demo", "size": 14, "runtime": "fixture-runtime",
                   "trained": "fixture build", "news": "fixture release", "features": [], "older": []}
        release_file = self.root / "release.json"
        release_file.write_text(json.dumps(release))
        options = {"RELEASE": str(release_file), "BASE": "/pet/"}
        self.run_script("site/build.sh", output, **options)
        for name in published:
            self.assertEqual(self.tree(output / name), published[name])
        config = self.config((output / "index.html").read_text())
        self.assertEqual(config["modelId"], "demo")
        (output / "old-entry.txt").write_text("old output must survive failures\n")
        before = self.tree(output)
        siblings = set(self.root.iterdir())
        failures = [
            ("invalid analytics", json.dumps(release), {"CF_ANALYTICS_TOKEN": "invalid"}),
            ("malformed release", "{", {}),
            ("unfilled render", json.dumps({**release, "trained": "@UNFILLED@"}), {}),
            ("invalid rendered config", json.dumps(release), {"MODEL_ID": 'bad"id'}),
        ]
        for name, content, overrides in failures:
            with self.subTest(name=name):
                release_file.write_text(content)
                self.run_script("site/build.sh", output, success=False, **options, **overrides)
                self.assertEqual(self.tree(output), before)
                self.assertEqual(self.config((output / "index.html").read_text()), config)
                self.assertEqual(set(self.root.iterdir()), siblings)
        release_file.write_text(json.dumps({**release, "news": "rebuilt release"}))
        self.run_script("site/build.sh", output, **options)
        self.assertFalse((output / "old-entry.txt").exists())
        self.assertEqual(self.config((output / "index.html").read_text())["news"], "rebuilt release")
        for name in published:
            self.assertEqual(self.tree(output / name), published[name])

    def test_publish_runtime_failure_allows_same_id_retry(self):
        model = self.model("model")
        runtime = self.runtime("runtime", "A")
        root = self.root / "published"
        options = {"STATIC_ROOT": str(root), "BASE": "pet", "FILE_MODE": "0640", "DIR_MODE": "0750"}
        # Exercise both a missing runtime and a copy that fails after copying files.
        broken = self.runtime("broken-runtime", "broken")
        (broken / "z-directory").mkdir()
        for source in (self.root / "missing-runtime", broken):
            with self.subTest(runtime=source.name):
                self.run_script("train/publish.sh", model, "demo", source, "fixture-runtime",
                                success=False, **options)
                self.assertEqual(list((root / "pet/models").iterdir()), [])
                self.assertEqual(list((root / "pet/runtime").iterdir()), [])
        self.run_script("train/publish.sh", model, "demo", runtime, "fixture-runtime", **options)
        published_model = root / "pet/models/demo"
        published_runtime = root / "pet/runtime/fixture-runtime"
        self.assertEqual(self.tree(published_model), self.tree(model))
        payload = self.tree(published_runtime)
        payload.pop("SHA256SUMS")
        self.assertEqual(payload, self.tree(runtime))
        for path in (published_model, published_runtime):
            self.assert_manifest(path)
            for entry in (path, *path.rglob("*")):
                self.assertEqual(entry.stat().st_mode & 0o777, 0o750 if entry.is_dir() else 0o640)

    def test_published_model_and_existing_runtime_are_immutable(self):
        model_a = self.model("model-a", b"model A\n")
        model_b = self.model("model-b", b"model B\n")
        runtime_a = self.runtime("runtime-a", "A")
        runtime_b = self.runtime("runtime-b", "B")
        root = self.root / "published"
        options = {"STATIC_ROOT": str(root), "BASE": "pet"}
        self.run_script("train/publish.sh", model_a, "demo", runtime_a, "fixture-runtime", **options)
        before = self.tree(root)
        self.run_script("train/publish.sh", model_b, "demo", runtime_b, "other-runtime",
                        success=False, **options)
        self.assertEqual(self.tree(root), before)
        self.run_script("train/publish.sh", model_b, "next", runtime_b, "fixture-runtime", **options)
        self.assertEqual(self.tree(root / "pet/models/demo"), self.tree(model_a))
        self.assertEqual(self.tree(root / "pet/models/next"), self.tree(model_b))
        original_runtime = {key.removeprefix("pet/runtime/fixture-runtime/"): value
                            for key, value in before.items() if key.startswith("pet/runtime/fixture-runtime/")}
        self.assertEqual(self.tree(root / "pet/runtime/fixture-runtime"), original_runtime)
        self.assert_manifest(root / "pet/runtime/fixture-runtime")

    def test_full_preview_refreshes_runtime_but_rebuild_does_not(self):
        model = self.model("model")
        runtime_a = self.runtime("runtime-a", "A")
        (runtime_a / "obsolete.wasm").write_bytes(b"obsolete")
        runtime_b = self.runtime("runtime-b", "B")
        (runtime_b / ".fixture").write_text("hidden runtime metadata\n")
        (runtime_b / "nested").mkdir()
        (runtime_b / "nested/fixture.txt").write_text("nested runtime metadata\n")
        output = self.root / "preview"
        selected = output / "pet/runtime/runtime"
        with self.preview(output, model, runtime_a) as url:
            body, headers = self.fetch(url + "/pet/runtime/runtime/transformers.min.js")
            self.assertEqual(body, (runtime_a / "transformers.min.js").read_bytes())
            self.assertEqual(headers["Cross-Origin-Opener-Policy"], "same-origin")
            self.assertEqual(headers["Cross-Origin-Embedder-Policy"], "require-corp")
        self.run_script("test/dev-site.sh", output, "--rebuild", MODEL=str(model), RUNTIME=str(runtime_b),
                        BASE="pet", MODEL_ID="demo")
        self.assertEqual(self.tree(selected), self.tree(runtime_a))
        with self.preview(output, model, runtime_b) as url:
            body, _ = self.fetch(url + "/pet/runtime/runtime/transformers.min.js")
            self.assertEqual(body, (runtime_b / "transformers.min.js").read_bytes())
            self.assertEqual(self.tree(selected), self.tree(runtime_b))
            with self.assertRaises(HTTPError) as error:
                self.fetch(url + "/pet/runtime/runtime/obsolete.wasm")
            with error.exception:
                self.assertEqual(error.exception.code, 404)
        previous = self.tree(selected)
        siblings = set(selected.parent.iterdir())
        self.run_script("test/dev-site.sh", output, success=False, MODEL=str(model),
                        RUNTIME=str(self.root / "missing-runtime"), BASE="pet", MODEL_ID="demo")
        self.assertEqual(self.tree(selected), previous)
        self.assertEqual(set(selected.parent.iterdir()), siblings)


if __name__ == "__main__":
    unittest.main()