use crate::models::ocr::model_metadata::ModelMetadata;
use crate::models::loader::{ModelLoader, ModelSession, ModelType};
use anyhow::Context;
use anyhow::Result;
use std::path::Path;
use tract_onnx::prelude::{tvec, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp};

pub struct OcrSession {
    pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
    pub model_metadata: ModelMetadata,
}
impl ModelSession for OcrSession {
    fn get_model_type(&self) -> ModelType {
        todo!("使用thiserror作为错误处理的库,thiserror 专门用于开发库（Library）");
    }
    fn desc(&self) -> String {
        "Ocr Model 加载成功".to_string()
    }
}
impl OcrSession {
    pub fn new<P>(model_path: P, model_metadata: ModelMetadata) -> Result<Self, anyhow::Error>
    where
        P: AsRef<Path>,
    {
        let session = ModelLoader::model_for_path(model_path)?.session;
        Ok(Self {
            session,
            model_metadata,
        })
    }

    pub fn model_from_bytes(
        model_bytes: &[u8],
        model_metadata: ModelMetadata,
    ) -> Result<Self, anyhow::Error> {
        let session = ModelLoader::model_from_bytes(model_bytes)?.session;
        Ok(Self {
            session,
            model_metadata,
        })
    }
    /// 对应 Python 的 _inference
    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())
    }
}