#pragma once
#include "inference/backend_config.hpp"

#include <onnxruntime/onnxruntime_cxx_api.h>
#include <filesystem>
#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>

// Private to the ORT backend (backends/ort_backend.cpp). Detects the best
// available ORT execution provider and applies it to a SessionOptions.
// Priority order: TensorRT EP > CUDA > ROCm > CPU.
//
// Note: "TensorRT" here is ORT's TensorRT *execution provider*, distinct from
// the raw-TensorRT backend (backends/trt_backend.cpp). This header never reaches
// core translation units.
//
// Detection is conservative: GetAvailableProviders() confirms ORT was compiled
// with the provider, then AppendExecutionProvider_* is attempted inside a
// try/catch so a missing runtime library degrades gracefully to the next tier.

enum class OrtProvider { CPU, CUDA, ROCm, TensorRT };

inline OrtProvider detect_ort_provider() {
    // ORT returns these in its own preference order (TensorRT, CUDA, ..., CPU
    // last), so the first recognised entry is the best available and the loop
    // returns on it.
    auto available = Ort::GetAvailableProviders();
    for (const auto& p : available) {
        // Only the TensorRT *EP* is a build-time opt-in — it needs the headers
        // and the profile plumbing below. CUDA is not: it is a plain ORT
        // provider, and gating its detection on the TRT flag (as this did) made
        // the CUDA branch unreachable in every build that did not also ask for
        // TensorRT. The symptom is silent rather than loud — inference simply
        // runs on the CPU and everything still returns correct answers — which
        // is why it survived: a 300-actor VR-012 grid cell took 76 s on the CPU
        // with the GPU idle at 212 MiB.
#ifdef SAE_ORT_WITH_TRT_EP
        if (p == "TensorrtExecutionProvider") return OrtProvider::TensorRT;
#endif
        if (p == "CUDAExecutionProvider")     return OrtProvider::CUDA;
        if (p == "ROCMExecutionProvider")     return OrtProvider::ROCm;
    }
    return OrtProvider::CPU;
}

inline const char* provider_name(OrtProvider p) {
    switch (p) {
        case OrtProvider::TensorRT: return "TensorRT-EP";
        case OrtProvider::CUDA:     return "CUDA";
        case OrtProvider::ROCm:     return "ROCm";
        default:                    return "CPU";
    }
}

// If trt_cfg.ort_cache_dir is set, configure ORT to write/read a pre-optimized
// .ort model for model_path. Must be called before AppendExecutionProvider_*.
inline void apply_ort_model_cache(Ort::SessionOptions& opts,
                                  const std::string& model_path,
                                  const BackendConfig& trt_cfg) {
    if (trt_cfg.ort_cache_dir.empty()) return;
    std::filesystem::create_directories(trt_cfg.ort_cache_dir);
    const std::string stem =
        std::filesystem::path(model_path).stem().string();
    const std::string cache_path =
        trt_cfg.ort_cache_dir + "/" + stem + ".ort";
    opts.SetOptimizedModelFilePath(cache_path.c_str());
}

// Apply the given provider to opts. Falls back to CPU on failure and returns the
// provider that was actually applied.
inline OrtProvider apply_ort_provider(Ort::SessionOptions& opts,
                                      OrtProvider provider,
                                      const char* label,
                                      const BackendConfig& trt_cfg = {}) {
#ifdef SAE_ORT_WITH_TRT_EP
    if (provider == OrtProvider::TensorRT) {
        try {
            std::filesystem::create_directories(trt_cfg.cache_dir);

            std::unordered_map<std::string, std::string> kv = {
                {"device_id",               "0"},
                {"trt_max_workspace_size",  "2147483648"},
                {"trt_fp16_enable",         trt_cfg.fp16 ? "1" : "0"},
                {"trt_int8_enable",         trt_cfg.int8 ? "1" : "0"},
                {"trt_engine_cache_enable", "1"},
                {"trt_engine_cache_path",   trt_cfg.cache_dir},
            };

            if (!trt_cfg.input_name.empty() && !trt_cfg.profile_min.empty()) {
                kv["trt_profile_min_shapes"] =
                    trt_cfg.input_name + ":" + trt_cfg.profile_min;
                kv["trt_profile_opt_shapes"] =
                    trt_cfg.input_name + ":" + trt_cfg.profile_opt;
                kv["trt_profile_max_shapes"] =
                    trt_cfg.input_name + ":" + trt_cfg.profile_max;
            }

            Ort::TensorRTProviderOptions trt_v2;
            trt_v2.Update(kv);
            opts.AppendExecutionProvider_TensorRT_V2(*trt_v2);

            std::cerr << "[" << label << "] TensorRT EP"
                      << (trt_cfg.fp16 ? " FP16" : "")
                      << (trt_cfg.int8 ? " INT8" : "")
                      << "  cache=" << trt_cfg.cache_dir
                      << (trt_cfg.profile_min.empty() ? "" :
                          "  profile=" + trt_cfg.profile_min
                          + "/" + trt_cfg.profile_opt
                          + "/" + trt_cfg.profile_max)
                      << "\n";
            return OrtProvider::TensorRT;
        } catch (const Ort::Exception& e) {
            std::cerr << "[" << label << "] TensorRT EP unavailable ("
                      << e.what() << "), trying CUDA\n";
            provider = OrtProvider::CUDA;
        }
    }
#endif  // SAE_ORT_WITH_TRT_EP
    if (provider == OrtProvider::CUDA) {
        try {
            OrtCUDAProviderOptions cuda{};
            cuda.device_id = 0;
            opts.AppendExecutionProvider_CUDA(cuda);
            std::cerr << "[" << label << "] CUDA provider\n";
            return OrtProvider::CUDA;
        } catch (const Ort::Exception& e) {
            std::cerr << "[" << label << "] CUDA unavailable ("
                      << e.what() << "), trying ROCm\n";
            provider = OrtProvider::ROCm;
        }
    }
    if (provider == OrtProvider::ROCm) {
        try {
            OrtROCMProviderOptions rocm{};
            rocm.device_id = 0;
            opts.AppendExecutionProvider_ROCM(rocm);
            std::cerr << "[" << label << "] ROCm provider\n";
            return OrtProvider::ROCm;
        } catch (const Ort::Exception& e) {
            std::cerr << "[" << label << "] ROCm unavailable ("
                      << e.what() << "), falling back to CPU\n";
        }
    }
    std::cerr << "[" << label << "] CPU provider\n";
    return OrtProvider::CPU;
}