#[cfg(feature = "ort-cuda")]
use ort::execution_providers::CUDAExecutionProvider;
#[cfg(feature = "ort-coreml")]
use ort::execution_providers::CoreMLExecutionProvider;
#[cfg(feature = "ort-directml")]
use ort::execution_providers::DirectMLExecutionProvider;
#[cfg(feature = "ort-openvino")]
use ort::execution_providers::OpenVINOExecutionProvider;
#[cfg(feature = "ort-tvm")]
use ort::execution_providers::TVMExecutionProvider;
#[cfg(feature = "ort-tensorrt")]
use ort::execution_providers::TensorRTExecutionProvider;
use ort::execution_providers::{CPUExecutionProvider, ExecutionProviderDispatch};

/// Supported execution providers for ONNX Runtime
#[derive(Debug, Clone)]
pub enum ExecutionProvider {
    /// CPU execution provider (always available)
    CPU,
    /// CoreML execution provider (macOS only)
    CoreML,
    /// CUDA execution provider (requires cuda feature)
    CUDA,
    /// TensorRT execution provider (requires tensorrt feature)
    TensorRT,
    /// TVM execution provider (requires tvm feature)
    TVM,
    /// OpenVINO execution provider (requires openvino feature)
    OpenVINO,
    /// DirectML execution provider (Windows only, requires directml feature)
    DirectML,
}

impl std::fmt::Display for ExecutionProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ExecutionProvider::CPU => write!(f, "CPU"),
            ExecutionProvider::CoreML => write!(f, "CoreML"),
            ExecutionProvider::CUDA => write!(f, "CUDA"),
            ExecutionProvider::TensorRT => write!(f, "TensorRT"),
            ExecutionProvider::TVM => write!(f, "TVM"),
            ExecutionProvider::OpenVINO => write!(f, "OpenVINO"),
            ExecutionProvider::DirectML => write!(f, "DirectML"),
        }
    }
}

impl std::str::FromStr for ExecutionProvider {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "cpu" => Ok(ExecutionProvider::CPU),
            "coreml" => Ok(ExecutionProvider::CoreML),
            "cuda" => Ok(ExecutionProvider::CUDA),
            "tensorrt" => Ok(ExecutionProvider::TensorRT),
            "tvm" => Ok(ExecutionProvider::TVM),
            "openvino" => Ok(ExecutionProvider::OpenVINO),
            "directml" => Ok(ExecutionProvider::DirectML),
            _ => Err(format!("Unknown execution provider: {}", s)),
        }
    }
}

impl ExecutionProvider {
    /// Returns all available execution providers for the current platform and features
    pub fn available_providers() -> Vec<ExecutionProvider> {
        vec![
            ExecutionProvider::CPU,
            #[cfg(all(target_os = "macos", feature = "ort-coreml"))]
            ExecutionProvider::CoreML,
            #[cfg(feature = "ort-cuda")]
            ExecutionProvider::CUDA,
            #[cfg(feature = "ort-tensorrt")]
            ExecutionProvider::TensorRT,
            #[cfg(feature = "ort-tvm")]
            ExecutionProvider::TVM,
            #[cfg(feature = "ort-openvino")]
            ExecutionProvider::OpenVINO,
            #[cfg(all(target_os = "windows", feature = "ort-directml"))]
            ExecutionProvider::DirectML,
        ]
    }

    /// Check if this execution provider is available on the current platform
    pub fn is_available(&self) -> bool {
        match self {
            ExecutionProvider::CPU => true,
            ExecutionProvider::CoreML => cfg!(target_os = "macos") && cfg!(feature = "ort-coreml"),
            ExecutionProvider::CUDA => cfg!(feature = "ort-cuda"),
            ExecutionProvider::TensorRT => cfg!(feature = "ort-tensorrt"),
            ExecutionProvider::TVM => cfg!(feature = "ort-tvm"),
            ExecutionProvider::OpenVINO => cfg!(feature = "ort-openvino"),
            ExecutionProvider::DirectML => {
                cfg!(target_os = "windows") && cfg!(feature = "ort-directml")
            }
        }
    }
}

impl ExecutionProvider {
    pub fn to_dispatch(&self) -> Option<ExecutionProviderDispatch> {
        match self {
            ExecutionProvider::CPU => Some(CPUExecutionProvider::default().build()),
            ExecutionProvider::CoreML => {
                #[cfg(target_os = "macos")]
                {
                    #[cfg(feature = "ort-coreml")]
                    {
                        use tap::Tap;

                        Some(
                            CoreMLExecutionProvider::default()
                                .with_model_format(
                                    ort::execution_providers::coreml::CoreMLModelFormat::MLProgram,
                                )
                                .build(),
                        )
                    }
                    #[cfg(not(feature = "ort-coreml"))]
                    {
                        tracing::error!("coreml support not compiled in");
                        None
                    }
                }
                #[cfg(not(target_os = "macos"))]
                {
                    tracing::error!("CoreML is only available on macOS");
                    None
                }
            }
            ExecutionProvider::CUDA => {
                #[cfg(feature = "ort-cuda")]
                {
                    Some(CUDAExecutionProvider::default().build())
                }
                #[cfg(not(feature = "ort-cuda"))]
                {
                    tracing::error!("CUDA support not compiled in");
                    None
                }
            }
            ExecutionProvider::TensorRT => {
                #[cfg(feature = "ort-tensorrt")]
                {
                    Some(TensorRTExecutionProvider::default().build())
                }
                #[cfg(not(feature = "ort-tensorrt"))]
                {
                    tracing::error!("TensorRT support not compiled in");
                    None
                }
            }
            ExecutionProvider::TVM => {
                #[cfg(feature = "ort-tvm")]
                {
                    Some(TVMExecutionProvider::default().build())
                }
                #[cfg(not(feature = "ort-tvm"))]
                {
                    tracing::error!("TVM support not compiled in");
                    None
                }
            }
            ExecutionProvider::OpenVINO => {
                #[cfg(feature = "ort-openvino")]
                {
                    Some(OpenVINOExecutionProvider::default().build())
                }
                #[cfg(not(feature = "ort-openvino"))]
                {
                    tracing::error!("OpenVINO support not compiled in");
                    None
                }
            }
            ExecutionProvider::DirectML => {
                #[cfg(target_os = "windows")]
                {
                    #[cfg(feature = "ort-directml")]
                    {
                        Some(DirectMLExecutionProvider::default().build())
                    }
                    #[cfg(not(feature = "ort-directml"))]
                    {
                        tracing::error!("DirectML support not compiled in");
                        None
                    }
                }
                #[cfg(not(target_os = "windows"))]
                {
                    tracing::error!("DirectML is only available on Windows");
                    None
                }
            }
        }
    }
}
