stock-rust 0.1.0

use anyhow::{anyhow, Result};

use crate::stocks::base::{COMMON_HK, COMMON_SH, COMMON_SZ, COMMON_US};

use super::base::{ApiCodeTransform, CommonCodeTransform};

const XUEQIU_SZ: &str = "SZ";
const XUEQIU_SH: &str = "SH";
const XUEQIU_HK: &str = "HK";
const XUEQIU_US: &str = "";

pub struct XueqiuCommonCodeTransform;

impl CommonCodeTransform for XueqiuCommonCodeTransform {
    fn transform(&self, code: &str) -> Result<String> {
        if code.starts_with(COMMON_SH) {
            return self.sh_transform(code);
        }
        if code.starts_with(COMMON_SZ) {
            return self.sz_transform(code);
        }
        if code.starts_with(COMMON_HK) {
            return self.hk_transform(code);
        }
        if code.starts_with(COMMON_US) {
            return self.us_transform(code);
        }

        Err(anyhow!("请检查统一代码是否正确"))
    }

    fn sz_transform(&self, code: &str) -> Result<String> {
        if !code.contains(COMMON_SZ) {
            return Err(anyhow!("请检查统一代码是否正确"));
        }
        Ok(format!("{}{}", XUEQIU_SZ, code.replacen(COMMON_SZ, "", 1)))
    }

    fn sh_transform(&self, code: &str) -> Result<String> {
        if !code.contains(COMMON_SH) {
            return Err(anyhow!("请检查统一代码是否正确"));
        }
        Ok(format!("{}{}", XUEQIU_SH, code.replacen(COMMON_SH, "", 1)))
    }

    fn hk_transform(&self, code: &str) -> Result<String> {
        if !code.contains(COMMON_HK) {
            return Err(anyhow!("请检查统一代码是否正确"));
        }
        Ok(format!("{}{}", XUEQIU_HK, code.replacen(COMMON_HK, "", 1)))
    }

    fn us_transform(&self, code: &str) -> Result<String> {
        if !code.contains(COMMON_US) {
            return Err(anyhow!("请检查统一代码是否正确"));
        }
        Ok(format!("{}{}", XUEQIU_US, code.replacen(COMMON_US, "", 1)))
    }
}

pub struct XueqiuApiCodeTransform;

impl ApiCodeTransform for XueqiuApiCodeTransform {
    fn transform(&self, code: &str) -> Result<String> {
        if code.contains(XUEQIU_SZ) {
            return self.sz_transform(code);
        }
        if code.contains(XUEQIU_SH) {
            return self.sh_transform(code);
        }
        if code.contains(XUEQIU_HK) {
            return self.hk_transform(code);
        }

        self.us_transform(code)
    }

    fn sz_transform(&self, code: &str) -> Result<String> {
        if !code.contains(XUEQIU_SZ) {
            return Err(anyhow!("请检查股票代码是否正确"));
        }
        Ok(format!("{}{}", COMMON_SZ, code.replacen(XUEQIU_SZ, "", 1)))
    }

    fn sh_transform(&self, code: &str) -> Result<String> {
        if !code.contains(XUEQIU_SH) {
            return Err(anyhow!("请检查股票代码是否正确"));
        }
        Ok(format!("{}{}", COMMON_SH, code.replacen(XUEQIU_SH, "", 1)))
    }

    fn hk_transform(&self, code: &str) -> Result<String> {
        if !code.contains(XUEQIU_HK) {
            return Err(anyhow!("请检查股票代码是否正确"));
        }
        Ok(format!("{}{}", COMMON_HK, code.replacen(XUEQIU_HK, "", 1)))
    }

    fn us_transform(&self, code: &str) -> Result<String> {
        if code.contains(XUEQIU_SZ) || code.contains(XUEQIU_SH) || code.contains(XUEQIU_HK) {
            return Err(anyhow!("请检查股票代码是否正确"));
        }
        Ok(format!("{}{}", COMMON_US, code.replacen(XUEQIU_US, "", 1)))
    }
}

pub struct XueqiuSearchCodeTransform;

impl XueqiuSearchCodeTransform {
    pub fn transform(&self, code: &str) -> Vec<String> {
        if let Some(raw) = code.strip_prefix("SZ") {
            return vec![format!("{}{}", COMMON_SZ, raw)];
        }
        if let Some(raw) = code.strip_prefix("SH") {
            return vec![format!("{}{}", COMMON_SH, raw)];
        }

        vec![format!("{}{}", COMMON_HK, code), format!("{}{}", COMMON_US, code)]
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_common_code_transform() {
        let t = XueqiuCommonCodeTransform;
        assert_eq!(t.transform("SH510500").unwrap(), "SH510500");
        assert_eq!(t.transform("SZ399001").unwrap(), "SZ399001");
        assert_eq!(t.transform("HKHSI").unwrap(), "HKHSI");
        assert_eq!(t.transform("USDJI").unwrap(), "DJI");
    }

    #[test]
    fn test_api_code_transform() {
        let t = XueqiuApiCodeTransform;
        assert_eq!(t.transform("SH510500").unwrap(), "SH510500");
        assert_eq!(t.transform("SZ399001").unwrap(), "SZ399001");
        assert_eq!(t.transform("HKHSI").unwrap(), "HKHSI");
        assert_eq!(t.transform("DJI").unwrap(), "USDJI");
    }

    #[test]
    fn test_search_code_transform() {
        let t = XueqiuSearchCodeTransform;
        assert_eq!(t.transform("SZ399001"), vec!["SZ399001"]);
        assert_eq!(t.transform("SH510500"), vec!["SH510500"]);
        assert_eq!(t.transform("HSI"), vec!["HKHSI", "USHSI"]);
    }
}