"""Install and launch the pinned Swift MLX runtime using an existing local model."""

import argparse
import hashlib
import inspect
import json
import os
import platform
import shlex
import shutil
import subprocess
import sys
import urllib.request
import venv
from pathlib import Path

BASE = "c69d1288440a0dc4e6401fc417098b07598dccd5"
REVISIONS = {
    4: "730aab9b0395b26f7d9cf4b0dfa6a4c788aff6fd",
    5: "8aff72b145212e62c15146e41dcbef35ea5fa9ff",
}
PATCHES = {
    "swift15-mlx-lm.patch": "f6f1d0bdafa45863bfbf93dac0398c481c993ea04fdf38b9bae98c643f89eaec",
    "enable-5bit.patch": "b985961eac3035e05ca4c9f3a8b283c26e6dd69bab4d113997fc04fe2a4f99cc",
    "swift15-server-cache.patch": "a7fbc0f0524ee7864d9f41a98a2e35acf0e63f53e369ee5b9a87ca12926a9eab",
}
MODEL_HASHES = {
    4: "f559e4559946bfad4ce20b93906028449b76e007f42f36bb83d2c2c5f6cb95a9",
    5: "ecda666d6d5f1059388a2ce4267d9ff8625398bb925d6c0840d0059ce17b676f",
}
SERVER_HASH = "791f8c1eb3431c4d24b4f9084f67b665c38e3d547ee4a307e52c5b2b6f12aba0"


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def local_config(model):
    config = json.loads((model / "config.json").read_text())
    quant = config.get("quantization", {})
    bits = quant.get("bits")
    if bits not in REVISIONS or quant != {
        "bits": bits,
        "group_size": 64,
        "mode": "affine",
    }:
        raise ValueError(
            "Expected the Swift 1.5 affine 4-bit or 5-bit / group-size-64 checkpoint."
        )
    if config.get("language_model_only") is not False:
        raise ValueError("Expected the complete Swift architecture configuration.")
    for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"):
        if not (model / name).is_file():
            raise ValueError(f"Missing local model asset: {name}")
    index = json.loads((model / "model.safetensors.index.json").read_text())
    shards = set(index["weight_map"].values())
    if not shards:
        raise ValueError("The checkpoint index is empty.")
    for name in shards:
        if Path(name).name != name or not name.endswith(".safetensors"):
            raise ValueError("The checkpoint index contains an invalid shard path.")
        if not (model / name).is_file():
            raise ValueError(
                f"Missing local weight shard: {name}; complete the existing download first."
            )
    return config, bits


def runtime_check(config, bits):
    from importlib.metadata import version

    from mlx_lm import server
    from mlx_lm.utils import _get_classes

    cls, args = _get_classes(config)
    args.from_dict(config)
    if cls.__module__ != "mlx_lm.models.qwen3_5_full":
        raise ValueError(
            "Swift architecture patch is missing. Use the serve launcher created by setup."
        )
    if digest(inspect.getsourcefile(cls)) != MODEL_HASHES[bits]:
        raise ValueError(
            "The installed architecture differs from the pinned Swift implementation."
        )
    if digest(server.__file__) != SERVER_HASH:
        raise ValueError("The tested server cache patch is missing or has changed.")
    versions = {
        name: version(name)
        for name in ("mlx", "mlx-metal", "mlx-lm", "transformers", "huggingface_hub")
    }
    expected = {
        "mlx": "0.32.2",
        "mlx-metal": "0.32.2",
        "mlx-lm": "0.32.0",
        "transformers": "5.14.1",
        "huggingface_hub": "1.31.0",
    }
    if versions != expected:
        raise ValueError(
            f"Installed package versions differ from the tested runtime: {versions}"
        )
    print(
        json.dumps(
            {
                "status": "PASS_RUNTIME_IDENTITY",
                "python": sys.executable,
                "model_class": cls.__module__ + "." + cls.__name__,
                "server": server.__file__,
                "versions": versions,
                "weights_loaded": False,
            }
        ),
        flush=True,
    )


def setup(model, bits, destination):
    if shutil.which("git") is None:
        raise ValueError(
            "Git is required. Install the macOS command-line tools, then retry."
        )
    if destination.exists():
        raise ValueError(
            f"Runtime directory already exists: {destination}. Use its serve launcher, or choose a new --runtime-dir."
        )
    destination.mkdir(parents=True)
    source = destination / "mlx-lm"
    envdir = destination / ".venv"
    python = envdir / "bin/python"
    patchdir = destination / "patches"
    patchdir.mkdir()
    patch_names = ["swift15-mlx-lm.patch"]
    if bits == 5:
        patch_names.append("enable-5bit.patch")
    patch_names.append("swift15-server-cache.patch")
    for name in patch_names:
        url = f"https://huggingface.co/ukisai/Swift-1.5-{bits}bit-MLX/resolve/{REVISIONS[bits]}/compatibility/{name}"
        with urllib.request.urlopen(url, timeout=60) as response:
            data = response.read()
        if hashlib.sha256(data).hexdigest() != PATCHES[name]:
            raise ValueError(f"Patch checksum mismatch: {name}")
        (patchdir / name).write_bytes(data)
    subprocess.run(["git", "init", "--quiet", str(source)], check=True)
    git = ["git", "-C", str(source)]
    subprocess.run(
        git + ["remote", "add", "origin", "https://github.com/ml-explore/mlx-lm.git"],
        check=True,
    )
    subprocess.run(git + ["fetch", "--depth", "1", "origin", BASE], check=True)
    subprocess.run(git + ["checkout", "--detach", BASE], check=True)
    for name in patch_names:
        subprocess.run(git + ["apply", "--check", str(patchdir / name)], check=True)
        subprocess.run(git + ["apply", str(patchdir / name)], check=True)
    venv.EnvBuilder(with_pip=True).create(envdir)
    subprocess.run(
        [
            str(python),
            "-I",
            "-m",
            "pip",
            "install",
            "mlx==0.32.2",
            "mlx-metal==0.32.2",
            "transformers==5.14.1",
            "huggingface_hub==1.31.0",
            "pillow==12.3.0",
            "safetensors==0.8.0",
            "-e",
            str(source),
        ],
        check=True,
    )
    subprocess.run([str(python), "-I", "-m", "pip", "check"], check=True)
    installed = destination / "swift_runtime.py"
    shutil.copyfile(__file__, installed)
    command = [str(python), "-I", str(installed), "start", "--model", str(model)]
    subprocess.run(command + ["--check-only"], check=True)
    launcher = destination / "serve"
    launcher.write_text("#!/bin/sh\nexec " + shlex.join(command) + ' "$@"\n')
    launcher.chmod(0o755)
    print(
        f"\nSetup complete. Start or restart with this exact command:\n{shlex.quote(str(launcher))}"
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="action", required=True)
    install = sub.add_parser(
        "setup",
        help="Install into a new isolated directory; reuse existing model files.",
    )
    install.add_argument("--model", required=True, type=Path)
    install.add_argument("--runtime-dir", type=Path)
    start = sub.add_parser(
        "start", help="Verify the selected Python runtime and serve a local checkpoint."
    )
    start.add_argument("--model", required=True, type=Path)
    start.add_argument("--check-only", action="store_true")
    args, extra = parser.parse_known_args()
    if args.action == "setup" and extra:
        parser.error("Unknown setup arguments: " + " ".join(extra))
    if sys.version_info[:2] != (3, 12):
        parser.error(
            "Use python3.12 for this tested runtime. Homebrew's Python 3.14 environment is separate."
        )
    if platform.system() != "Darwin" or platform.machine() != "arm64":
        parser.error("This installer and launcher target native Apple Silicon macOS.")
    model = args.model.expanduser().resolve()
    try:
        config, bits = local_config(model)
        if args.action == "setup":
            target = (
                args.runtime_dir
                or Path.home() / ".local/share/swift15-mlx" / f"{bits}bit"
            )
            setup(model, bits, target.expanduser().resolve())
            return
        os.environ["HF_HUB_OFFLINE"] = "1"
        os.environ["TRANSFORMERS_OFFLINE"] = "1"
        runtime_check(config, bits)
        if args.check_only:
            return
        os.execv(
            sys.executable,
            [
                sys.executable,
                "-I",
                "-m",
                "mlx_lm.server",
                "--model",
                str(model),
                "--host",
                "127.0.0.1",
                "--port",
                "8080",
                *extra,
            ],
        )
    except (
        OSError,
        ValueError,
        KeyError,
        ImportError,
        subprocess.CalledProcessError,
    ) as exc:
        raise SystemExit(f"Swift runtime setup/start failed: {exc}") from None


if __name__ == "__main__":
    main()