use anyhow::{anyhow, Result};
use serde::Deserialize;
use std::borrow::Cow;
use std::collections::HashMap;

// ==========================================
// 3. 字符集核心结构体 (重命名为 Charset)
// ==========================================
#[derive(Debug, Clone)]
pub struct Charset {
    // 使用 Cow 统一静态切片和动态读取的 Vec<String>，内部实现真正的零拷贝
    pub tokens: Vec<Cow<'static, str>>,
    // 反向查找表，保证字符转索引为 O(1)
    pub char_to_idx: HashMap<Cow<'static, str>, usize>,
    // 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速)
    // pub valid_indices: HashSet<usize>,
}

impl Charset {
    // 内部底层统一收拢构造
    pub fn new(tokens: Vec<Cow<'static, str>>) -> Self {
        let mut char_to_idx = HashMap::with_capacity(tokens.len());
        for (idx, token) in tokens.iter().enumerate() {
            char_to_idx.entry(token.clone()).or_insert(idx);
            // 如果字符集有重复，保留第一个遇到的索引 (符合 Python .index 逻辑)
            // char_to_idx.entry(token.to_string()).or_insert(idx);
        }

        Self {
            tokens,
            char_to_idx,
        }
    }

    // --- 业务策略方法 ---

    /// 将字符转为索引，不存在返回 -1 (保持与原 Python 库行为一致)
    pub fn char_to_index(&self, char_str: &str) -> i32 {
        if let Some(&idx) = self.char_to_idx.get(char_str) {
            idx as i32
        } else {
            -1
        }
    }

    /// 将索引转为字符引用，零拷贝。若越界返回 None
    pub fn index_to_char_ref(&self, index: usize) -> Option<&str> {
        self.tokens.get(index).map(|cow| cow.as_ref())
    }

    pub fn is_valid_char(&self, char_str: &str) -> bool {
        self.char_to_idx.get(char_str).is_some()
    }
    pub fn size(&self) -> usize {
        self.tokens.len()
    }
}

// ==========================================
// 4. 标准 Display 接口实现 (对应 __str__)
// ==========================================
impl std::fmt::Display for Charset {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Charset [Total Size: {}", self.size(),)
    }
}



// =====================================================================
// 1. 辅助定义的枚举与结构体
// =====================================================================

#[derive(Debug, Clone, Copy, Deserialize)]
#[serde(rename_all = "snake_case")] // 支持 json 中写 "zero_to_one" 或 "minus_one_to_one"
pub enum Normalization {
    /// 映射到 [0.0, 1.0] -> pixel / 255.0
    ZeroToOne,
    /// 映射到 [-1.0, 1.0] -> (pixel / 255.0 - 0.5) / 0.5
    MinusOneToOne,
}

impl Normalization {
    /// 统一归一化计算逻辑
    #[inline(always)]
    pub fn normalize(&self, pixel: f32) -> f32 {
        match self {
            Normalization::ZeroToOne => pixel / 255.0,
            Normalization::MinusOneToOne => (pixel / 255.0 - 0.5) / 0.5,
        }
    }
}

/// 图像缩放策略枚举
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Resize {
    /// 固定宽高，例如 (64, 64)
    Fixed(u32, u32),
    /// 高度固定，宽度根据原始比例动态计算（对应 Python 的 [-1, H]）
    DynamicWidth(u32),
    /// 单字识别的正方形切图（对应 Python 的 word 为 True 且 [-1, H]）
    Square(u32),
}

/// 仅用于反序列化 JSON 的中间临时结构体（DTO）
#[derive(Deserialize)]
struct ModelMetadataDto {
    charset: Vec<String>,
    word: bool,
    #[serde(alias = "image")]
    resize: Vec<i32>,
    channel: u8,
    /// 新增：允许在配置文件中指定归一化策略。
    /// 使用 serde(default) 可以在不配置时提供一个默认值（比如默认 ZeroToOne）
    #[serde(default = "default_normalization")]
    normalization: Normalization,
}
fn default_normalization() -> Normalization {
    Normalization::ZeroToOne
}

#[derive(Debug, Clone)]
pub struct ModelMetadata {
    /// 字符集管理器
    pub charset: Charset,
    /// 是否为单字识别模型
    pub word: bool,
    /// 预处理的缩放策略
    pub resize: Resize,
    /// 图像通道数 (1 或 3)
    pub channel: u8,
    /// 新增：传递给核心业务使用的归一化配置
    pub normalization: Normalization,
}

impl ModelMetadata {
    // --- 优雅的工厂模式构造器 ---
    /// 通用的静态切片转换构造器
    pub fn from_static_slice(
        slice: &[&'static str],
        word: bool,
        resize: Resize,
        channel: u8,
        normalization: Normalization,
    ) -> Self {
        let tokens: Vec<Cow<'static, str>> = slice.iter().map(|&s| Cow::Borrowed(s)).collect();
        Self {
            charset: Charset::new(tokens),
            word,
            resize,
            channel,
            normalization,
        }
    }
    pub fn from_json_str(json_str: &str) -> Result<Self> {
        let dto: ModelMetadataDto = serde_json::from_str(json_str)
            .map_err(|e| anyhow!("JSON 反序列化失败，请检查字段是否完整: {}", e))?;

        // 1. 将 DTO 的字符串数组转化为强类型的 Charset
        let tokens: Vec<Cow<'static, str>> =
            dto.charset.into_iter().map(|s| Cow::Owned(s)).collect();
        let charset = Charset::new(tokens);

        // 2. 解析 resize 策略（重现 Python 的复杂条件判断）
        if dto.resize.len() != 2 {
            return Err(anyhow!(
                "'resize (or image)' 字段必须是包含两个元素的数组，例如 [-1, 64]"
            ));
        }
        let r0 = dto.resize[0];
        let r1 = dto.resize[1];

        let resize = if r0 == -1 {
            if dto.word {
                // 如果 word 为 true，且包含 -1，Python 里是 resize 为 (r1, r1) 的正方形
                Resize::Square(r1 as u32)
            } else {
                // 如果 word 为 false，且包含 -1，Python 里是高度固定为 r1，宽度按原图比例缩放
                Resize::DynamicWidth(r1 as u32)
            }
        } else {
            // 正常的固定宽高
            Resize::Fixed(r0 as u32, r1 as u32)
        };

        Ok(Self {
            charset,
            word: dto.word,
            resize,
            channel: dto.channel,
            normalization: dto.normalization,
        })
    }
    /// 机制 2：从内存字节流加载（极大地方便 include_bytes! 或网络下载）
    pub fn from_json_bytes(bytes: &[u8]) -> Result<Self> {
        let json_str = std::str::from_utf8(bytes)
            .map_err(|e| anyhow!("JSON 字节流不是合法的 UTF-8 编码: {}", e))?;
        Self::from_json_str(json_str)
    }
}