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 SINA_SZ: &str = "sz";
const SINA_SH: &str = "sh";
const SINA_HK: &str = "hk";
const SINA_US: &str = "gb_";

pub struct SinaCommonCodeTransform;

impl CommonCodeTransform for SinaCommonCodeTransform {
    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.starts_with(COMMON_SZ) {
            return Err(anyhow!("请检查统一代码是否正确"));
        }
        Ok(format!("{}{}", SINA_SZ, code.replacen(COMMON_SZ, "", 1)))
    }

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

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

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

pub struct SinaApiCodeTransform;

impl ApiCodeTransform for SinaApiCodeTransform {
    fn transform(&self, code: &str) -> Result<String> {
        if code.starts_with(SINA_SZ) {
            return self.sz_transform(code);
        }
        if code.starts_with(SINA_SH) {
            return self.sh_transform(code);
        }
        if code.starts_with(SINA_HK) {
            return self.hk_transform(code);
        }
        if code.starts_with(SINA_US) {
            return self.us_transform(code);
        }

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

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

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

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

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

pub struct SinaSearchCodeTransform;

impl SinaSearchCodeTransform {
    pub fn transform(&self, code: &str) -> Vec<String> {
        if let Some(raw) = code.strip_prefix("us") {
            return vec![format!("{}{}", COMMON_US, raw)];
        }
        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)];
        }
        if let Some(raw) = code.strip_prefix("hk") {
            return vec![format!("{}{}", COMMON_HK, raw)];
        }
        if let Some(raw) = code.strip_prefix("of") {
            return vec![format!("{}{}", COMMON_SZ, raw), format!("{}{}", COMMON_SH, raw)];
        }

        Vec::new()
    }
}

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

    #[test]
    fn test_common_code_transform() {
        let t = SinaCommonCodeTransform;
        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(), "gb_dji");
    }

    #[test]
    fn test_api_code_transform() {
        let t = SinaApiCodeTransform;
        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("gb_dji").unwrap(), "USdji");
    }

    #[test]
    fn test_search_code_transform() {
        let t = SinaSearchCodeTransform;
        assert_eq!(t.transform("usdji"), vec!["USdji"]);
        assert_eq!(t.transform("sz399001"), vec!["SZ399001"]);
        assert_eq!(t.transform("sh510500"), vec!["SH510500"]);
        assert_eq!(t.transform("hkHSI"), vec!["HKHSI"]);
        assert_eq!(t.transform("of510500"), vec!["SZ510500", "SH510500"]);
    }
}