"""Merge-update `$OUT_DIR/run_identity.json` so plugins can record backend facts.

Shell writes a stub before pytest; plugins add device_name / timer / vendor
without clobbering earlier keys. Missing OUT_DIR is a silent no-op so unit
tests and ad-hoc pytest invocations do not fail.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Any, Dict, Optional


IDENTITY_NAME = "run_identity.json"


def identity_path(out_dir: Optional[str] = None) -> Optional[Path]:
    root = out_dir or os.environ.get("OUT_DIR", "").strip()
    if not root:
        return None
    return Path(root) / IDENTITY_NAME


def load_identity(out_dir: Optional[str] = None) -> Dict[str, Any]:
    path = identity_path(out_dir)
    if path is None or not path.is_file():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}
    return data if isinstance(data, dict) else {}


def update_identity(out_dir: Optional[str] = None, **fields: Any) -> Dict[str, Any]:
    path = identity_path(out_dir)
    if path is None:
        return dict(fields)
    path.parent.mkdir(parents=True, exist_ok=True)
    current = load_identity(out_dir)
    for key, value in fields.items():
        if value is None:
            continue
        current[key] = value
    tmp = path.with_name(path.name + ".tmp")
    tmp.write_text(json.dumps(current, indent=2, ensure_ascii=False) + "\n",
                   encoding="utf-8")
    os.replace(tmp, path)
    return current


def init_identity(out_dir: str, **fields: Any) -> Dict[str, Any]:
    payload = {"schema_version": 1, "provider": "flaggems"}
    payload.update(fields)
    os.environ["OUT_DIR"] = str(out_dir)
    return update_identity(out_dir, **payload)


def main(argv=None) -> int:
    p = argparse.ArgumentParser()
    p.add_argument("out_dir")
    p.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"))
    p.add_argument("--op", default=os.environ.get("OP", ""))
    p.add_argument("--op-file", default=os.environ.get("OP_FILE", ""))
    p.add_argument("--flaggems-dir", default=os.environ.get("FLAGGEMS_DIR", ""))
    p.add_argument("--autotune-mode", default="")
    p.add_argument("--device", default=os.environ.get("ZL_DEVICE", ""))
    p.add_argument("--test-file", default=os.environ.get("TEST_FILE", ""))
    p.add_argument("--shape-mode", default=os.environ.get("SHAPE_MODE", ""))
    p.add_argument("--bench-level", default=os.environ.get("BENCH_LEVEL", "core"))
    # "isolated" or "shared(<path>)". Recorded because a shared libtuner cache
    # lets whichever side runs first hand its winners to the other, and after
    # the fact the two runs are otherwise indistinguishable.
    p.add_argument("--tuner-cache", default="")
    args = p.parse_args(argv)
    init_identity(
        args.out_dir,
        backend=args.backend,
        op=args.op,
        op_file=args.op_file,
        flaggems_dir=args.flaggems_dir,
        autotune_mode=args.autotune_mode,
        device=args.device,
        test_file=args.test_file,
        shape_mode=args.shape_mode,
        bench_level=args.bench_level,
        tuner_cache=args.tuner_cache,
    )
    return 0


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