#!/usr/bin/env python3
"""Create a pinned Qwen checkpoint metadata attestation, locally and once.

No model runtime is imported, no weight contents are read, no socket is opened,
and an existing output is never overwritten.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from qwen_runtime.launcher import LauncherError, build_checkpoint_attestation  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model-path",
        type=Path,
        default=Path(os.environ.get("QWEN_MODEL_PATH", os.environ.get("MODEL_PATH", ""))),
        required=not bool(os.environ.get("QWEN_MODEL_PATH") or os.environ.get("MODEL_PATH")),
        help="absolute local Qwen checkpoint directory",
    )
    parser.add_argument(
        "--model-id",
        default=os.environ.get("QWEN_EXACT_MODEL_ID", os.environ.get("MODEL_ID", "")),
        required=not bool(os.environ.get("QWEN_EXACT_MODEL_ID") or os.environ.get("MODEL_ID")),
    )
    parser.add_argument(
        "--revision",
        default=os.environ.get("QWEN_EXACT_MODEL_REVISION", os.environ.get("MODEL_REVISION", "")),
        required=not bool(os.environ.get("QWEN_EXACT_MODEL_REVISION") or os.environ.get("MODEL_REVISION")),
        help="immutable 40-character checkpoint commit",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="defaults to <model-path>/.omni-checkpoint.json; existing files are refused",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        model_root = args.model_path.resolve(strict=True)
        attestation = build_checkpoint_attestation(
            model_root,
            model_id=args.model_id,
            revision=args.revision,
        )
        output = (args.output or model_root / ".omni-checkpoint.json").absolute()
        output_parent = output.parent.resolve(strict=True)
        output = output_parent / output.name
        if output_parent != model_root:
            raise LauncherError("attestation output must be directly inside the checkpoint directory")
        payload = (json.dumps(attestation, ensure_ascii=False, indent=2) + "\n").encode("utf-8")
        descriptor = os.open(output, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        try:
            with os.fdopen(descriptor, "wb") as handle:
                handle.write(payload)
                handle.flush()
                os.fsync(handle.fileno())
        except Exception:
            try:
                output.unlink(missing_ok=True)
            finally:
                raise
    except (LauncherError, OSError) as exc:
        print(json.dumps({"ok": False, "error": str(exc)}, ensure_ascii=False), file=sys.stderr)
        return 2
    print(
        json.dumps(
            {
                "ok": True,
                "path": str(output),
                "sha256": hashlib.sha256(payload).hexdigest(),
                "model_id": attestation["model_id"],
                "revision": attestation["revision"],
                "weight_shards": attestation["expected_shard_count"],
                "model_started": False,
                "network_used": False,
            },
            ensure_ascii=False,
            sort_keys=True,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())