use crate::loader::Error;
use crate::loader::error::{BuildError, ParseError, Result};
use crate::types::Session;
use ddddocr_core::traits::Loader;
use ort::session::Session as OrtSession;
use ort::session::builder::SessionBuilder;
use std::sync::{Arc, Mutex};

// pub struct OrtModelLoader;
//
// impl OrtModelLoader {
//     /// 获取针对 ORT 后端的链式构建器
//     pub fn builder() -> OrtModelBuilder {
//         OrtModelBuilder::default()
//     }
// }
/// ORT 专用的链式构建器
#[derive(Debug, Clone)]
pub struct ModelLoader {
    use_gpu: bool,
    device_id: i32,
    intra_threads: Option<usize>,
}
impl Default for ModelLoader {
    fn default() -> Self {
        Self {
            use_gpu: false,
            device_id: 0,
            intra_threads: None,
        }
    }
}
impl ModelLoader {
    /// 开启或关闭 GPU 加速
    pub fn use_gpu(mut self, enable: bool) -> Self {
        self.use_gpu = enable;
        self
    }

    /// 指定 GPU 设备 ID
    pub fn device_id(mut self, id: i32) -> Self {
        self.device_id = id;
        self
    }
    pub fn num_threads(mut self, threads: usize) -> Self {
        self.intra_threads = Some(threads);
        self
    }

    /// 内部辅助方法：根据当前的配置构建 ORT 底层的 SessionBuilder
    fn create_session_builder(&self) -> Result<SessionBuilder> {
        let mut builder = OrtSession::builder().map_err(|e| BuildError::BuildFailed(e))?;
        // 如果用户显式设置了线程数，则配置给 ORT
        if let Some(threads) = self.intra_threads {
            builder = builder
                .with_intra_threads(threads)
                .map_err(|e| BuildError::Threads(format!("设置线程数失败: {e}")))?;
        }

        if self.use_gpu {
            // 根据 ort 库版本配置 CUDA 执行提供者 (Execution Provider)
            #[cfg(feature = "cuda")]
            {
                use ort::ep::CUDAExecutionProvider;
                let cuda_ep = CUDAExecutionProvider::default().with_device_id(self.device_id);
                builder = builder
                    .with_execution_providers([cuda_ep.build()])
                    .map_err(|e| {
                        BuildError::EnabledCudaFailed(format!("配置 CUDA 硬件加速失败: {e}"))
                    })?
            }
            #[cfg(not(feature = "cuda"))]
            {
                // 如果用户明确开启了 GPU，但 Feature 没编译进去，明确抛错提醒
                return Err(BuildError::NotEnabledCuda(
                    "未启用 CUDA 支持：请在 Cargo.toml 中为 ddddocr-ort 开启 `cuda` feature"
                        .to_string(),
                ))?;
            }
        }

        Ok(builder)
    }
}

impl Loader for ModelLoader {
    type Session = Session;
    type Error = Error;

    fn build_for_path<P>(&self, model_path: P) -> Result<Session>
    where
        P: AsRef<std::path::Path>,
    {
        let path_ref = model_path.as_ref();
        let mut builder = self.create_session_builder()?;
        // Session::builder() 会返回 Result<SessionBuilder, ort::Error>

        let session = builder
            .commit_from_file(path_ref)
            .map_err(|e| ParseError::Path(path_ref.display().to_string(), e))?;

        Ok(Arc::new(Mutex::new(session))) // 这里的session需要包装下
    }
    /// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用）
    fn build_from_bytes(&self, model_bytes: &[u8]) -> Result<Session> {
        let mut builder = self.create_session_builder()?;

        let session = builder
            .commit_from_memory(model_bytes)
            .map_err(ParseError::Bytes)?;

        Ok(Arc::new(Mutex::new(session)))
    }
}