#!python
# SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC
#
# SPDX-License-Identifier: Apache-2.0

"""Run tt-metal's triage utility from a source or wheel installation."""

from importlib.util import find_spec
import os
from pathlib import Path
import runpy
import sys


def _reexec_in_container_triage_environment() -> None:
    toolchain_directory = os.environ.get("TTLANG_TOOLCHAIN_DIR", "").strip()
    if not toolchain_directory:
        return

    triage_environment = Path(toolchain_directory) / "tt-triage-venv"
    if Path(sys.prefix).resolve() == triage_environment.resolve():
        return

    triage_python = triage_environment / "bin" / "python"
    if not triage_python.is_file():
        return

    # Device-management packages in the container have independent version
    # constraints, so triage runs with its isolated dependencies.
    environment = os.environ.copy()
    environment["VIRTUAL_ENV"] = str(triage_environment)
    existing_path = environment.get("PATH", "")
    environment["PATH"] = str(triage_environment / "bin") + (
        f"{os.pathsep}{existing_path}" if existing_path else ""
    )
    os.execve(
        str(triage_python),
        [str(triage_python), str(Path(__file__).resolve()), *sys.argv[1:]],
        environment,
    )


def _find_triage_directory() -> Path:
    runtime_root = (
        os.environ.get("TT_METAL_RUNTIME_ROOT", "").strip()
        or os.environ.get("TT_METAL_HOME", "").strip()
    )
    if runtime_root:
        triage_directory = Path(runtime_root) / "tools" / "triage"
        if (triage_directory / "triage.py").is_file():
            return triage_directory

    toolchain_directory = os.environ.get("TTLANG_TOOLCHAIN_DIR", "").strip()
    if toolchain_directory:
        triage_directory = Path(toolchain_directory) / "tt-metal" / "tools" / "triage"
        if (triage_directory / "triage.py").is_file():
            return triage_directory

    triage_spec = find_spec("triage")
    if triage_spec is not None and triage_spec.submodule_search_locations:
        for package_directory in triage_spec.submodule_search_locations:
            triage_directory = Path(package_directory)
            if (triage_directory / "triage.py").is_file():
                return triage_directory

    raise SystemExit(
        "tt-triage could not find tools/triage in the installed tt-metal runtime"
    )


def main() -> None:
    _reexec_in_container_triage_environment()
    triage_directory = _find_triage_directory()
    sys.path.insert(0, str(triage_directory))
    runpy.run_path(str(triage_directory / "triage.py"), run_name="__main__")


if __name__ == "__main__":
    main()
