scuriolus 0.2.0

use chrono::{DateTime, Utc};
use std::marker::PhantomData;
#[cfg(test)]
use strum::IntoEnumIterator;
use surrealdb::engine::local::{Db, RocksDb};
use surrealdb::Surreal;

use crate::core::{CoreError, CoreResult};
#[cfg(test)]
use crate::data::{Crypto, KlineInterval};
use crate::market::{Kline, KlinesParams, Market};

const DB_ADDRESS_BASE: &str = "db/kline-store/";

const DB_NAMESPACE: &str = "scuriolus";
const DB_NAME: &str = "main";

#[derive(Debug, serde::Serialize, serde::Deserialize)]
struct SectionKlines {
    #[serde(with = "chrono::serde::ts_seconds")]
    pub begin: DateTime<Utc>,
    #[serde(with = "chrono::serde::ts_seconds")]
    pub end: DateTime<Utc>,
    pub valid: bool,
}

#[derive(Debug, serde::Serialize)]
struct KlineWithId {
    #[allow(dead_code)]
    id: String,
    kline: Kline,
}

#[derive(Debug, serde::Deserialize)]
struct KlineWithoutId {
    kline: Kline,
}

impl From<Kline> for KlineWithId {
    fn from(kline: Kline) -> Self {
        KlineWithId {
            id: kline.open_time.to_string(),
            kline,
        }
    }
}

impl From<KlineWithoutId> for Kline {
    fn from(kline: KlineWithoutId) -> Self {
        kline.kline
    }
}

/// A [`Kline`] store to prevent too much exchanges klines sources. Use an embeded database.
#[derive(Debug)]
pub struct KlineStore<M: Market> {
    db: Surreal<Db>,
    marker: PhantomData<M>,
}

impl<M: Market> Clone for KlineStore<M> {
    fn clone(&self) -> Self {
        Self {
            db: self.db.clone(),
            marker: PhantomData,
        }
    }
}

impl<M: Market> KlineStore<M> {
    /// Create a new [`KlineStore`]. Should only called once per [`Market`].
    pub async fn new() -> Self {
        Self::true_new("main").await
    }

    // Used to test with different klin_store, avoiding the same db
    #[cfg(test)]
    pub async fn test_new(name: &str) -> Self {
        Self::true_new(name).await
    }

    #[cfg(test)]
    pub async fn clear(&self, asset: Crypto, quote: Crypto) -> CoreResult<()> {
        for interval in KlineInterval::iter() {
            let (register, table) = get_register_and_table(&KlinesParams::new(
                asset,
                quote,
                interval,
                Utc::now(),
                Utc::now(),
            ));
            let _: Vec<KlineWithoutId> = self.db.delete(table).await?;
            let _: Vec<SectionKlines> = self.db.delete(register).await?;
        }
        Ok(())
    }

    async fn true_new(name: &str) -> Self {
        let db =
            Surreal::new::<RocksDb>(format!("{}/{}/{}", DB_ADDRESS_BASE, M::source_name(), name))
                .await
                .expect("error creating test with mexc db");
        db.use_ns(DB_NAMESPACE)
            .use_db(DB_NAME)
            .await
            .expect("error using test with mexc db");

        Self {
            db,
            marker: PhantomData,
        }
    }

    pub async fn inject_klines(&self, params: &KlinesParams, klines: Vec<Kline>) -> CoreResult<()> {
        if klines.is_empty() {
            tracing::warn!("No klines to inject");
            return Ok(());
        }
        //TODO checker les klines ? (continues ?)

        tracing::debug!("Injecting {} klines", klines.len());

        let (register, table) = get_register_and_table(params);

        let begin = klines.first().unwrap().open_time;
        let end = klines.last().unwrap().close_time;

        let mut response = self.db
            .query("SELECT * FROM type::table($register) WHERE begin <= $begin AND end >= $begin ORDER BY begin ASC")
            .query("SELECT * FROM type::table($register) WHERE begin <= $end AND end >= $end ORDER BY end DESC")
            .query("DELETE type::table($register) WHERE begin >= $begin AND begin <= $end OR end >= $begin AND end <= $end OR begin <= $begin AND end >= $end")
            .bind(("register", register.to_string()))
            .bind(("begin", begin.timestamp()))
            .bind(("end", end.timestamp()))
            .await?;

        let begin_klines: Vec<SectionKlines> = response.take(0)?;
        let end_klines: Vec<SectionKlines> = response.take(1)?;

        if begin_klines.len() > 1 || end_klines.len() > 1 {
            return Err(CoreError::db_error("Db corrupted"));
        }

        tracing::trace!(
            "begin_klines found locally: {:#?}, \nend_klines found locally: {:#?}",
            begin_klines,
            end_klines
        );

        let new_begin = match begin_klines.first() {
            Some(kline) => kline.begin,
            None => begin,
        };
        let new_end = match end_klines.first() {
            Some(kline) => kline.end,
            None => end,
        };

        let section = SectionKlines {
            begin: new_begin,
            end: new_end,
            valid: true,
        };

        let _: Option<SectionKlines> = self
            .db
            .insert((register, section.begin.timestamp()))
            .content(section)
            .await?;

        let klines_with_id: Vec<KlineWithId> = klines.into_iter().map(KlineWithId::from).collect();

        for kline in klines_with_id {
            let _: Vec<KlineWithoutId> = self.db.upsert(table.clone()).content(kline).await?;
        }

        Ok(())
    }

    pub async fn try_get_klines(&self, params: &KlinesParams) -> CoreResult<(Vec<Kline>, bool)> {
        let (register, table) = get_register_and_table(params);

        let start = params.start_time().timestamp();
        let end = params.end_time().timestamp();

        tracing::trace!(
            "SELECT * FROM {} WHERE begin <= {} AND end >= {} ({} -> {})",
            &register,
            start,
            end,
            params.start_time(),
            params.end_time()
        );

        let mut response = self
            .db
            .query("SELECT * FROM type::table($register) WHERE begin <= $start AND end >= $start")
            .bind(("register", register))
            .bind(("start", start))
            .await?;

        let content: Vec<SectionKlines> = response.take(0)?;

        if content.is_empty() {
            return Ok((Vec::new(), false));
        }

        let section = content.first().unwrap();

        let (last_kline_should_begin, last_kline_might_end) =
            params.interval().get_time_bounds(section.end)?;

        let complete = section.end.timestamp() == end
            || (last_kline_should_begin.timestamp() > end
                && end <= last_kline_might_end.timestamp());

        tracing::debug!("Asking DB for klines");
        tracing::trace!("table: {}, start: {}, end: {}", table, start, end);

        let mut kline_response = self
                .db
                .query("SELECT * FROM type::table($table) WHERE kline.open_time >= $start AND kline.close_time <= $end_section AND kline.close_time <= $end ORDER BY kline.open_time ASC")
                .bind(("start", start))
                .bind(("end_section", section.end.timestamp()))
                .bind(("end", end))
                .bind(("table", table))
                .await?;
        let klines: Vec<KlineWithoutId> = kline_response.take(0)?;
        Ok((klines.into_iter().map(Kline::from).collect(), complete))
    }
}

pub fn get_register_and_table(params: &KlinesParams) -> (String, String) {
    let interval = params.interval().short_string();
    (
        format!("{}{}_{}_reg", params.asset(), params.quote(), interval),
        format!("{}{}_{}", params.asset(), params.quote(), interval),
    )
}

#[cfg(test)]
mod tests {
    use chrono::{TimeZone as _, Utc};
    use rust_decimal::Decimal;

    use crate::{
        data::{Crypto, KlineInterval},
        market::{mock_context, Kline, Market as _, MockMarket},
    };

    use super::{super::super::market::KlinesParams, KlineStore};

    const ASSET: Crypto = Crypto::BTC;
    const QUOTE: Crypto = Crypto::USDT;

    fn kline_default() -> Kline {
        Kline {
            open_time: Utc.with_ymd_and_hms(2024, 2, 1, 0, 0, 0).unwrap(),
            open: Decimal::from(2),
            high: Decimal::from(5),
            low: Decimal::from(1),
            close: Decimal::from(3),
            volume: Decimal::from(50),
            close_time: Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap(),
            quote_asset_volume: Decimal::from(100),
        }
    }

    #[tokio::test]
    async fn simple() {
        let _ctx = mock_context();

        tracing::debug!("{}", MockMarket::source_name());

        let store = KlineStore::<MockMarket>::test_new("kline_store-simple").await;
        store.clear(ASSET, QUOTE).await.unwrap();

        let kline_default = kline_default();

        let jan = Kline {
            open_time: Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
            close_time: Utc.with_ymd_and_hms(2024, 2, 1, 0, 0, 0).unwrap(),
            ..kline_default
        };
        let feb = Kline {
            open_time: Utc.with_ymd_and_hms(2024, 2, 1, 0, 0, 0).unwrap(),
            close_time: Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap(),
            ..kline_default
        };
        let mar = Kline {
            open_time: Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap(),
            close_time: Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap(),
            ..kline_default
        };

        let params1 = &KlinesParams::new(
            ASSET,
            QUOTE,
            KlineInterval::OneMonth,
            Utc.with_ymd_and_hms(2024, 2, 1, 0, 0, 0).unwrap(),
            Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap(),
        );

        let klines1 = vec![feb.clone(), mar.clone()];

        store.inject_klines(params1, klines1).await.unwrap();

        let params2 = &KlinesParams::new(
            Crypto::BTC,
            Crypto::USDT,
            KlineInterval::OneMonth,
            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
            Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap(),
        );

        let klines2 = vec![jan.clone(), feb.clone()];

        store.inject_klines(params2, klines2).await.unwrap();

        let (klines, complete) = store
            .try_get_klines(&KlinesParams::new(
                ASSET,
                QUOTE,
                KlineInterval::OneMonth,
                Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
                Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap(),
            ))
            .await
            .unwrap();

        assert!(complete);
        assert_eq!(klines, vec![jan, feb, mar]);
    }

    #[tokio::test]
    async fn uncomplete() {
        let _ctx = mock_context();

        let store = KlineStore::<MockMarket>::test_new("kline_store-uncomplete").await;
        store.clear(ASSET, QUOTE).await.unwrap();

        let kline_default = kline_default();

        let feb = Kline {
            open_time: Utc.with_ymd_and_hms(2024, 2, 1, 0, 0, 0).unwrap(),
            close_time: Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap(),
            ..kline_default
        };
        let mar = Kline {
            open_time: Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap(),
            close_time: Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap(),
            ..kline_default
        };

        let params = &KlinesParams::new(
            ASSET,
            QUOTE,
            KlineInterval::OneMonth,
            Utc.with_ymd_and_hms(2024, 2, 1, 0, 0, 0).unwrap(),
            Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap(),
        );

        let klines = vec![feb.clone(), mar.clone()];

        store.inject_klines(params, klines).await.unwrap();

        let (klines, complete) = store
            .try_get_klines(&KlinesParams::new(
                ASSET,
                QUOTE,
                KlineInterval::OneMonth,
                Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
                Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap(),
            ))
            .await
            .unwrap();

        assert!(!complete);
        assert_eq!(klines, vec![]);

        let (klines, complete) = store
            .try_get_klines(&KlinesParams::new(
                ASSET,
                QUOTE,
                KlineInterval::OneMonth,
                Utc.with_ymd_and_hms(2024, 2, 1, 0, 0, 0).unwrap(),
                Utc.with_ymd_and_hms(2024, 5, 1, 0, 0, 0).unwrap(),
            ))
            .await
            .unwrap();

        assert!(!complete);
        assert_eq!(klines, vec![feb, mar]);
    }
}