use anyhow::Context;
use std::io::Cursor;
use tract_onnx::onnx;
use tract_onnx::prelude::*;

/// OCR 模型：包含路径和字符集

pub enum ModelType {
    Ocr,
    Det,
    Custom,
}
// 定义统一的 trait
pub trait ModelSession {
    fn get_model_type(&self) -> ModelType;
    fn desc(&self) -> String;
}

pub struct ModelLoader {
    pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
}

impl ModelLoader {
    pub fn model_for_path<P>(model_path: P) -> anyhow::Result<Self>
    where
        P: AsRef<std::path::Path>,
    {
        let session = onnx()
            .model_for_path(model_path)
            .with_context(|| "加载 ONNX 模型失败，请检查路径是否正确")?
            .into_optimized()?
            .into_runnable()?;
        Ok(Self { session })
    }
    /// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用）
    pub fn model_from_bytes(model_bytes: &[u8]) -> anyhow::Result<Self> {
        // 使用 std::io::Cursor 将 &[u8] 包装为可读的流（实现 std::io::Read）
        let mut cursor = Cursor::new(model_bytes);

        let session = onnx()
            .model_for_read(&mut cursor)
            .with_context(|| "从内存字节流解析 ONNX 模型失败")?
            .into_optimized()?
            .into_runnable()?;

        Ok(Self { session })
    }
}