"""Create an InferenceSession that survives a provider this machine cannot run.

Shipped inside the onnxruntime package by scripts/runtime/stage_python_bundle.ps1
so that it is importable wherever `import onnxruntime` works:

    from onnxruntime.uort import create_session

    sess = create_session("model.onnx")            # best available provider
    sess = create_session("model.onnx", prefer=["DmlExecutionProvider"])

Two ONNX Runtime behaviours make the obvious call unreliable.

1. Provider fallback is narrower than it looks. The python binding retries with
   the remaining providers only for exceptions derived from ValueError or
   RuntimeError (onnxruntime/capi/onnxruntime_inference_collection.py, "Falling
   back to ... and retrying"). TensorRT fails with
   `onnxruntime.capi.onnxruntime_pybind11_state.Fail`, whose base class is
   Exception, so an engine that cannot be built for this GPU aborts session
   creation even when CUDAExecutionProvider is in the list.

2. DirectML refuses to share a session: "DML EP can only be used with CPU EPs".
   Passing ['CUDAExecutionProvider', 'DmlExecutionProvider', ...] fails for
   DirectML *and* makes the failure look like a CUDA/TensorRT error, because the
   provider that is initialised first reports it.

So the helper tries whole provider lists, best first, and keeps DirectML in a
list of its own. This is what the C++ UORT runtime does by probing providers
before creating a session.
"""

from __future__ import annotations

# Order used when no provider is requested: TensorRT partitions the graph and
# CUDA runs the rest, which is the fastest combination on a Turing-or-newer GPU.
DEFAULT_PREFERENCE = [
    "TensorrtExecutionProvider",
    "CUDAExecutionProvider",
    "DmlExecutionProvider",
    "OpenVINOExecutionProvider",
    "CPUExecutionProvider",
]


def _candidate_lists(prefer, available):
    names = [name for name in prefer if name in available]
    if "CPUExecutionProvider" not in names:
        names.append("CPUExecutionProvider")

    lists = []
    has_trt_cuda = "TensorrtExecutionProvider" in names and "CUDAExecutionProvider" in names
    if has_trt_cuda:
        lists.append(["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"])

    for name in names:
        if name == "CPUExecutionProvider":
            continue
        # Already covered by the TensorRT + CUDA combination above; CUDA on its
        # own is still worth trying, because the TensorRT engine may be the only
        # part that cannot run on this GPU.
        if has_trt_cuda and name == "TensorrtExecutionProvider":
            continue
        lists.append([name, "CPUExecutionProvider"])

    lists.append(["CPUExecutionProvider"])
    return lists


def create_session(model, prefer=None, provider_options=None, quiet=False, allow_fallback=True, **kwargs):
    """Return an InferenceSession that runs on the best provider this host has.

    model             path to the .onnx file
    prefer            provider names, best first (default: DEFAULT_PREFERENCE)
    provider_options  {provider_name: {option: value}}, applied to that provider
    quiet             do not report which provider was dropped and why
    allow_fallback    False = only try the first candidate (explicit --ep)
    kwargs            forwarded to ort.InferenceSession (sess_options, ...)

    Raises the last error when no provider can run the model.
    """
    import onnxruntime as ort

    options = provider_options or {}
    available = ort.get_available_providers()
    candidates = _candidate_lists(prefer or DEFAULT_PREFERENCE, available)
    if not allow_fallback:
        candidates = candidates[:1]

    last_error = None
    for names in candidates:
        providers = [(name, options[name]) if name in options else name for name in names]
        try:
            session = ort.InferenceSession(model, providers=providers, **kwargs)
            if last_error is not None and not quiet:
                print("using %s" % session.get_providers()[0])
            return session
        except Exception as exc:  # noqa: BLE001 - Fail is not a RuntimeError
            last_error = exc
            if not quiet:
                print("%s cannot run this model: %s" % ("+".join(names), exc))

    raise last_error