use crate::loader::ModelLoader;
use anyhow::Context;
use ddddocr_core::error::{DdddError, Result, TensorErrorReason};
use ddddocr_core::{DetEngine, DetOutput, InferenceEngine};
use ndarray::Ix3;
use std::path::Path;
use tract_onnx::prelude::{Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tvec};
#[derive(Debug)]
pub struct DetSession {
    pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
}

impl DetSession {
    pub fn new<P>(model_path: P) -> Result<Self>
    where
        P: AsRef<Path>,
    {
        let session = ModelLoader::model_for_path(&model_path)?.session;
        Ok(Self { session })
    }

    pub fn model_from_bytes(model_bytes: &[u8]) -> Result<Self> {
        let session = ModelLoader::model_from_bytes(model_bytes)?.session;
        Ok(Self { session })
    }
    // pub fn inference(&self, tensor: Tensor) -> anyhow::Result<Tensor> {
    //     // tract 的 run 会返回一个 Vec<TValue>，我们通常只需要第一个输出
    //     // let result = self.ocr.run(tvec!(tensor.into()))?;
    //     let mut result = self
    //         .session
    //         .run(tvec!(tensor.into()))
    //         .context("执行模型推理失败")?;
    //     println!("模型输出原始数据: {:?}", result);
    //     Ok(result.swap_remove(0).into_tensor())
    // }
}

impl InferenceEngine for DetSession {
    type Output = DetOutput; // 明确绑定 OCR 小枚举
    fn inference(&self, input_array: ndarray::Array4<f32>) -> Result<Self::Output> {
        // tract 的 run 会返回一个 Vec<TValue>，我们通常只需要第一个输出
        // let result = self.ocr.run(tvec!(tensor.into()))?;
        let tensor = Tensor::from(input_array);

        let mut result = self.session.run(tvec!(tensor.into())).map_err(|_| {
            DdddError::Inference(TensorErrorReason::EngineError(
                "执行模型推理失败".to_string(),
            ))
        })?;
        println!("模型输出原始数据: {:?}", result);
        // Ok(result.swap_remove(0).into_tensor())
        let raw_tensor = result.swap_remove(0).into_tensor();
        let array_d = raw_tensor.into_array::<f32>().map_err(|_| {
            DdddError::Inference(TensorErrorReason::EngineError(
                "Tract 实体张量无法转换为 ndarray::ArrayD".to_string(),
            ))
        })?;
        // 提前利用克隆(Clone)备份好当前未转维度前的真实 shape (Vec<usize>)
        let actual_shape = array_d.shape().to_vec();

        let array3 = array_d.into_dimensionality::<Ix3>().map_err(|_| {
            DdddError::Inference(TensorErrorReason::TensorDimensionMismatch {
                expected: "3D 检测矩阵 [Batch, Box_Count, Box_Attributes]".to_string(),
                actual: actual_shape, // 优雅降维失败时动态捕获
            })
        })?;
        Ok(DetOutput::Detection(array3))

        // 在引擎内部消化掉 DatumType 强耦合
    }
}

impl DetEngine for DetSession {}