stock-rust 0.1.0

use anyhow::Result;
use async_trait::async_trait;
use encoding_rs::GBK;

use crate::stocks::base::{dedup_codes, parse_num, percent};
use crate::stocks::transforms::base::CommonCodeTransform;
use crate::stocks::transforms::tencent::{TencentCommonCodeTransform, TencentSearchCodeTransform};
use crate::types::{Stock, StockApi};

#[derive(Clone)]
pub struct TencentApi {
    client: reqwest::Client,
}

impl TencentApi {
    pub fn new(client: reqwest::Client) -> Self {
        Self { client }
    }

    fn parse_line(code: &str, line: &str) -> Stock {
        let params_unformatted = line.split('=').nth(1).unwrap_or("").trim();
        let params: Vec<&str> = params_unformatted.trim_matches('"').split('~').collect();

        let now = parse_num(params.get(3).copied());
        let low = parse_num(params.get(34).copied());
        let high = parse_num(params.get(33).copied());
        let yesterday = parse_num(params.get(4).copied());

        Stock {
            code: code.to_uppercase(),
            name: params.get(1).unwrap_or(&"---").to_string(),
            now,
            low,
            high,
            yesterday,
            percent: percent(now, yesterday),
        }
    }
}

#[async_trait]
impl StockApi for TencentApi {
    async fn get_stock(&self, code: &str) -> Result<Stock> {
        let transformer = TencentCommonCodeTransform;
        let transformed = transformer.transform(code)?;
        let url = format!("https://qt.gtimg.cn/q={}", transformed);
        let body = self.client.get(url).send().await?.bytes().await?;
        let (text, _, _) = GBK.decode(&body);
        let line = text.split(";\n").next().unwrap_or_default();

        if !line.contains(&transformed) {
            return Ok(Stock::default_with_code(code));
        }

        Ok(Self::parse_line(code, line))
    }

    async fn get_stocks(&self, codes: &[String]) -> Result<Vec<Stock>> {
        let codes = dedup_codes(codes);
        if codes.is_empty() {
            return Ok(Vec::new());
        }

        let transformer = TencentCommonCodeTransform;
        let transformed = transformer.transforms(&codes)?;
        let url = format!("https://qt.gtimg.cn/q={}", transformed.join(","));
        let body = self.client.get(url).send().await?.bytes().await?;
        let (text, _, _) = GBK.decode(&body);
        let rows: Vec<&str> = text.split(";\n").collect();

        Ok(codes
            .iter()
            .enumerate()
            .map(|(index, code)| {
                let transformed = transformer.transform(code).unwrap_or_default();
                let row = rows.get(index).copied().unwrap_or_default();
                if !row.contains(&transformed) {
                    return Stock::default_with_code(code);
                }
                Self::parse_line(code, row)
            })
            .collect())
    }

    async fn search_stocks(&self, key: &str) -> Result<Vec<Stock>> {
        let url = format!(
            "https://smartbox.gtimg.cn/s3/?v=2&t=all&c=1&q={}",
            urlencoding::encode(key)
        );
        let body = self.client.get(url).send().await?.bytes().await?;
        let (text, _, _) = GBK.decode(&body);

        let sanitized = text.replace("v_hint=\"", "").replace('"', "");
        let rows: Vec<&str> = sanitized.split('^').collect();

        let search_transformer = TencentSearchCodeTransform;
        let codes: Vec<String> = rows
            .iter()
            .map(|row| {
                let mut sp = row.split('~');
                let tp = sp.next().unwrap_or_default();
                let code = sp.next().unwrap_or_default();
                search_transformer.transform(tp, code)
            })
            .filter(|s| !s.is_empty())
            .collect();

        self.get_stocks(&dedup_codes(&codes)).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::stocks::transforms::base::CommonCodeTransform;
    use crate::stocks::transforms::tencent::TencentCommonCodeTransform;

    #[test]
    fn test_to_api_code() {
        let transformer = TencentCommonCodeTransform;
        assert_eq!(transformer.transform("SH510500").unwrap(), "sh510500");
        assert_eq!(transformer.transform("SZ000001").unwrap(), "sz000001");
        assert_eq!(transformer.transform("HKhsi").unwrap(), "hkHSI");
        assert_eq!(transformer.transform("USDJI").unwrap(), "usDJI");
    }

    #[test]
    fn test_parse_line() {
        let line = "v_sh510500=\"1~500ETF~510500~7.224~7.149~7.147~2724781~1351032~1373749~7.224~63~7.223~329~7.222~370~7.221~286~7.220~143~7.225~3893~7.226~884~7.227~323~7.228~5728~7.229~190~~20200731154041~0.075~1.05~7.280~7.085~7.224/2724781/1964297124~2724781~196430~~~~7.280~7.085~2.73~~~0.00~7.864~6.434~1.20\"";
        let s = TencentApi::parse_line("SH510500", line);
        assert_eq!(s.code, "SH510500");
        assert_eq!(s.name, "500ETF");
        assert!((s.now - 7.224).abs() < 1e-12);
        assert!((s.low - 7.085).abs() < 1e-12);
        assert!((s.high - 7.28).abs() < 1e-12);
        assert!((s.yesterday - 7.149).abs() < 1e-12);
        assert!((s.percent - (7.224 / 7.149 - 1.0)).abs() < 1e-12);
    }
}