#!/usr/bin/env python3
"""Atomically update the Accessibility3R inference status file.

This helper deliberately uses only the Python standard library so it can run
inside the existing Accessibility3D environment without adding a dependency.
"""

from __future__ import annotations

import argparse
import json
import os
from datetime import datetime
from pathlib import Path


def now_iso() -> str:
    return datetime.now().astimezone().isoformat(timespec="seconds")


def load_existing(path: Path) -> dict:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return {}
    return value if isinstance(value, dict) else {}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--path", type=Path, required=True)
    parser.add_argument("--status", choices=("running", "completed", "failed"), required=True)
    parser.add_argument("--stage-key")
    parser.add_argument("--stage-label")
    parser.add_argument("--percent", type=float)
    parser.add_argument("--stage-end-percent", type=float)
    parser.add_argument("--stage-expected-seconds", type=float)
    parser.add_argument("--total-estimated-seconds", type=float)
    parser.add_argument("--message")
    parser.add_argument("--exit-code", type=int)
    args = parser.parse_args()

    path = args.path.expanduser().resolve()
    path.parent.mkdir(parents=True, exist_ok=True)
    current = load_existing(path)
    timestamp = now_iso()
    current.setdefault("started_at", timestamp)
    current["updated_at"] = timestamp
    current["status"] = args.status

    if args.stage_key is not None:
        current["stage_key"] = args.stage_key
        current["stage_started_at"] = timestamp
    if args.stage_label is not None:
        current["stage_label"] = args.stage_label
    if args.percent is not None:
        current["percent"] = max(0.0, min(100.0, args.percent))
    if args.stage_end_percent is not None:
        current["stage_end_percent"] = max(0.0, min(100.0, args.stage_end_percent))
    if args.stage_expected_seconds is not None:
        current["stage_expected_seconds"] = max(0.0, args.stage_expected_seconds)
    if args.total_estimated_seconds is not None:
        current["total_estimated_seconds"] = max(1.0, args.total_estimated_seconds)
    if args.message is not None:
        current["message"] = args.message
    if args.exit_code is not None:
        current["exit_code"] = args.exit_code

    if args.status == "completed":
        current.pop("failed_at", None)
        current.pop("exit_code", None)
        current["percent"] = 100.0
        current["stage_end_percent"] = 100.0
        current["completed_at"] = timestamp
    elif args.status == "failed":
        current.pop("completed_at", None)
        current["failed_at"] = timestamp
    else:
        current.pop("completed_at", None)
        current.pop("failed_at", None)
        current.pop("exit_code", None)

    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    temporary.write_text(
        json.dumps(current, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)
    return 0


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