#!/usr/bin/env python3
"""Capture exact available tool and Python package versions."""

from __future__ import annotations

import argparse
import importlib.metadata
import json
import platform
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any

from pipeline_common import REPO_ROOT, atomic_write_json, utc_now


TOOLS = ("git", "curl", "wget", "cmake", "ninja", "onnx-mlir", "onnx-mlir-opt", "mlir-opt", "flatc")
PACKAGES = (
    "numpy",
    "onnx",
    "onnxruntime",
    "jsonschema",
    "pytest",
    "pillow",
    "ai-edge-litert",
    "tensorflow",
    "tflite-runtime",
)


def command_version(executable: str) -> dict[str, Any]:
    path = shutil.which(executable)
    if path is None:
        return {"available": False, "path": None, "version_output": None, "exit_code": None}
    for flag in ("--version", "-version"):
        try:
            completed = subprocess.run([path, flag], capture_output=True, text=True, timeout=10, check=False)
        except (OSError, subprocess.TimeoutExpired):
            continue
        output = (completed.stdout + completed.stderr).strip()[:2000]
        if output or completed.returncode == 0:
            return {"available": True, "path": path, "version_output": output, "exit_code": completed.returncode}
    return {"available": True, "path": path, "version_output": "UNKNOWN", "exit_code": None}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=REPO_ROOT / "environment" / "tool_versions.json")
    args = parser.parse_args()
    packages: dict[str, Any] = {}
    for package in PACKAGES:
        try:
            packages[package] = {"available": True, "version": importlib.metadata.version(package)}
        except importlib.metadata.PackageNotFoundError:
            packages[package] = {"available": False, "version": None}
    result = {
        "captured_at": utc_now(),
        "python": {
            "version": platform.python_version(),
            "implementation": platform.python_implementation(),
            "executable": sys.executable,
        },
        "platform": platform.platform(),
        "tools": {tool: command_version(tool) for tool in TOOLS},
        "packages": packages,
    }
    atomic_write_json(args.output.resolve(), result)
    print(json.dumps(result, sort_keys=True))
    return 0


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