apple-quant-algorithmic 0.2.0

use std::{marker::PhantomData, range::Range};

use bevy::log::error;
use databento::{
	DateTimeLike,
	dbn::{FIXED_PRICE_SCALE, RecordRef, RecordRefEnum, Schema, Side, decode::AsyncDbnDecoder},
};
use smallstr::SmallString;
use time::{Duration, OffsetDateTime, UtcDateTime, UtcOffset};
use tokio::{fs::File, io::BufReader, time::sleep};

use crate::{
	aggregation::{Trade, TradeTradeTimestamp},
	backend::{DataBackend, OrderIdGenerator, OrdersBackend, RealtimeDataBackend},
	instrument::{InstrumentData, InstrumentSpec, InstrumentTicker, IsFloatingPoint, XSpec},
	order::{ActiveOrderGoals, ActiveStateGoal, DeferredOrderActions},
	order_manager::{OrderManager, OrdersCapacitySpec},
	price::{AbsolutePrice, FromCorrectedPrice},
	schema::SchemaFlags,
	strategy::Strategy,
	timestamp::{TickTimestamp, Timestamp, Timestamped, TradeTimestamp, TradeTimestamped},
	volume::{
		AggressiveVolume, AggressorSide, DirectionalExposure, DirectionlessVolume,
		FromCorrectedVolume,
	},
};

mod historical;
mod live;
mod symbology;

use historical::*;
use live::*;
use symbology::*;

impl From<&Schema> for SchemaFlags {
	fn from(value: &Schema) -> Self {
		match value {
			// Self::Mbo => SchemaFlags::AddCancelModifyFillClear | SchemaFlags::Trades,
			// Self::Mbp10 => SchemaFlags::LimitedAddCancelModifyFillClear | SchemaFlags::Trades,
			// Self::Mbp1 => SchemaFlags::LimitedAddCancelModifyFillClear | SchemaFlags::Trades,
			// Self::Tbbo => SchemaFlags::Trades,
			Schema::Trades => SchemaFlags::Trades,
			// Self::Ohlcv1S => SchemaFlags::OHLC1s,
			// Self::Ohlcv1M => SchemaFlags::OHLC1m,
			_ => SchemaFlags::empty(),
		}
	}
}

impl From<Schema> for SchemaFlags {
	fn from(value: Schema) -> Self {
		Self::from(&value)
	}
}

pub(crate) fn schema_value(schema: &Schema) -> u8 {
	match schema {
		Schema::Mbo => 8,
		Schema::Mbp10 => 7,
		Schema::Mbp1 => 6,
		Schema::Tbbo => 5,
		Schema::Tcbbo => 4,
		Schema::Trades => 3,
		Schema::Ohlcv1S => 2,
		Schema::Ohlcv1M => 1,
		_ => 0,
	}
}

impl DateTimeLike for Timestamp {
	fn to_date_time(self) -> OffsetDateTime {
		OffsetDateTime::from_unix_timestamp_nanos(self.as_utc_nanos()).unwrap()
	}
}

pub struct Databento<IS: InstrumentSpec + Send, OB: OrdersBackend<IS>> {
	key: Option<SmallString<[u8; 64]>>,

	realtime: Option<DatabentoLive>,
	historical: Option<DatabentoHistorical>,

	_is: PhantomData<IS>,
	_ob: PhantomData<OB>,
}

impl<IS: InstrumentSpec + Send, OB: OrdersBackend<IS>> Databento<IS, OB> {
	pub fn new(key: Option<&str>) -> Self {
		Self {
			key: key.map(|key| SmallString::from_str(key)),

			realtime: None,
			historical: None,

			_is: PhantomData::default(),
			_ob: PhantomData::default(),
		}
	}

	async fn new_realtime(
		&mut self,
		instrument_ticker: &InstrumentTicker,
		start_timestamp: impl DateTimeLike,
	) -> Result<&mut DatabentoLive, ()> {
		if let Some(mut realtime_databento) = self.realtime.take() {
			let _ = realtime_databento
				.client
				.close()
				.await;
		}

		let Some(key) = &self.key else {
			return Err(());
		};

		let Ok((realtime_databento, _schema_flags)) = DatabentoLive::new(
			instrument_ticker,
			start_timestamp,
			&[Schema::Trades],
			key.as_str(),
		)
		.await
		else {
			return Err(());
		};

		unsafe {
			self.realtime = Some(realtime_databento);

			Ok(self
				.realtime
				.as_mut()
				.unwrap_unchecked())
		}
	}

	async fn current_realtime(&mut self) -> Option<&mut DatabentoLive> {
		self.realtime.as_mut()
	}

	async fn initialize_historical(&mut self) -> Result<&mut DatabentoHistorical, ()> {
		let Some(key) = &self.key else {
			return Err(());
		};

		self.historical = Some(DatabentoHistorical::new(
			key.as_str(),
		));

		unsafe {
			Ok(self
				.historical
				.as_mut()
				.unwrap_unchecked())
		}
	}

	async fn historical(&mut self) -> Result<&mut DatabentoHistorical, ()> {
		if self.historical.is_none() {
			return self
				.initialize_historical()
				.await;
		}

		let Some(databento_historical) = self.historical.as_mut() else {
			return Err(());
		};

		Ok(databento_historical)
	}

	const HISTORICAL_LIVE_OVERLAP_SHORT: Duration = Duration::seconds(5);
	const HISTORICAL_LIVE_OVERLAP_LONG: Duration = Duration::minutes(10);

	async fn fetch_once(
		&mut self,
		instrument_ticker: &InstrumentTicker,
		utc_date_time_range: Range<UtcDateTime>,
	) -> impl IntoIterator<Item = TradeTradeTimestamp<IS>> {
		debug_assert!(utc_date_time_range.start <= utc_date_time_range.end);

		let now = UtcDateTime::now();
		let exact_cutoff = now - Duration::hours(24);

		let live_safe_cutoff = exact_cutoff + Self::HISTORICAL_LIVE_OVERLAP_SHORT;
		let historical_safe_cutoff = exact_cutoff - Self::HISTORICAL_LIVE_OVERLAP_SHORT;

		debug_assert!(utc_date_time_range.end < (now + Duration::seconds(5)));

		let mut market_data = Vec::with_capacity(100_000);

		// Historical client only.
		if utc_date_time_range.start <= historical_safe_cutoff
			&& utc_date_time_range.end <= historical_safe_cutoff
		{
			let historical_databento = self
				.historical()
				.await
				.unwrap();

			let mut stream = historical_databento
				.stream(
					instrument_ticker,
					utc_date_time_range,
				)
				.await
				.unwrap();

			while let Ok(Some(record_ref)) = stream
				.decode_record_ref()
				.await
			{
				let Some(trade_trade_timestamp) = process_record_trades(record_ref) else {
					continue;
				};

				market_data.push(trade_trade_timestamp);
			}

			return market_data;
		}

		// Live client only.
		if utc_date_time_range.start >= live_safe_cutoff
			&& utc_date_time_range.end >= live_safe_cutoff
		{
			let realtime_databento = self
				.new_realtime(
					instrument_ticker,
					utc_date_time_range
						.start
						.to_offset(UtcOffset::UTC),
				)
				.await
				.unwrap();

			while let Ok(Some(record_ref)) = realtime_databento
				.client
				.next_record()
				.await
			{
				let Some(trade_trade_timestamp) = process_record_trades(record_ref) else {
					continue;
				};

				let trade_date_time = trade_trade_timestamp
					.trade_timestamp()
					.to_date_time()
					.to_utc();

				if trade_date_time >= utc_date_time_range.end {
					break;
				}

				market_data.push(trade_trade_timestamp);
			}

			let _ = realtime_databento
				.client
				.close()
				.await;

			self.realtime = None;

			return market_data;
		}

		let historical_start = if utc_date_time_range.start < historical_safe_cutoff {
			utc_date_time_range.start
		} else {
			historical_safe_cutoff
		};

		let realtime_end = if utc_date_time_range.end > live_safe_cutoff {
			utc_date_time_range.end
		} else {
			live_safe_cutoff
		};

		let realtime_databento = self
			.new_realtime(
				instrument_ticker,
				OffsetDateTime::UNIX_EPOCH,
			)
			.await
			.unwrap();

		let mut realtime_satisfied_start = false;

		while let Ok(Some(record_ref)) = realtime_databento
			.client
			.next_record()
			.await
		{
			let Some(trade_trade_timestamp) = process_record_trades(record_ref) else {
				continue;
			};

			let trade_date_time = trade_trade_timestamp
				.trade_timestamp()
				.to_date_time()
				.to_utc();

			if trade_date_time >= realtime_end {
				break;
			}

			if trade_date_time < historical_start {
				realtime_satisfied_start = true;
				continue;
			}

			market_data.push(trade_trade_timestamp);
		}

		let _ = realtime_databento
			.client
			.close()
			.await;

		self.realtime = None;

		if realtime_satisfied_start {
			return market_data;
		}

		let first_live_timestamp = market_data
			.first()
			.unwrap()
			.trade_timestamp()
			.to_date_time()
			.to_utc();

		debug_assert!(first_live_timestamp < live_safe_cutoff);

		let historical_databento = self
			.historical()
			.await
			.unwrap();

		let start_utc = UtcDateTime::now();

		while UtcDateTime::now() - start_utc < Duration::seconds(30) {
			let available_end = historical_databento
				.available_end()
				.await
				.unwrap();

			if available_end > first_live_timestamp {
				break;
			}

			sleep(std::time::Duration::from_millis(100)).await;
		}

		let mut historical_market_data = Vec::with_capacity(10_000);

		let mut stream = historical_databento
			.stream(
				instrument_ticker,
				Range {
					start: historical_start,
					end: (first_live_timestamp + Duration::nanoseconds(1)),
				},
			)
			.await
			.unwrap();

		while let Ok(Some(record_ref)) = stream
			.decode_record_ref()
			.await
		{
			let Some(trade_trade_timestamp) = process_record_trades(record_ref) else {
				continue;
			};

			historical_market_data.push(trade_trade_timestamp);
		}

		let inclusive_start_live_idx = market_data
			.iter()
			.position(|trade_trade_timestamp| {
				trade_trade_timestamp
					.trade_timestamp()
					.to_date_time()
					.to_utc() > first_live_timestamp
			})
			.unwrap();

		historical_market_data.extend(
			market_data
				.into_iter()
				.skip(inclusive_start_live_idx),
		);

		historical_market_data
	}
}

impl<
	IS: InstrumentSpec + Send,
	OB: OrdersBackend<IS>,
	OrdersCS: OrdersCapacitySpec,
	S: Strategy<IS, OB, OrdersCS> + Send,
> RealtimeDataBackend<IS, OB, OrdersCS, S> for Databento<IS, OB>
{
	async fn initialize_realtime(
		&mut self,
		instrument_ticker: &InstrumentTicker,
		recent_trade_timestamp: Option<&TradeTimestamp>,
	) {
		let start_timestamp = recent_trade_timestamp
			.map(|trade_timestamp| trade_timestamp.to_date_time())
			.unwrap_or_else(|| OffsetDateTime::now_utc());

		let _ = self
			.new_realtime(
				instrument_ticker,
				start_timestamp,
			)
			.await;
	}

	async fn poll_realtime<'instrument_data, 'aggregated_data>(
		&mut self,
		instrument_data: &'instrument_data mut InstrumentData<
			'instrument_data,
			'aggregated_data,
			IS,
		>,
		strategy: &mut S,
		order_manager: &OrderManager<IS, OrdersCS>,
		directional_exposure: &DirectionalExposure<IS>,
		active_order_goals: &ActiveOrderGoals<IS>,
		active_state_goal: &ActiveStateGoal,
		deferred_order_actions: &mut DeferredOrderActions<IS>,
		order_id_generator: &mut OrderIdGenerator,
		recent_trade_timestamp: Option<&TradeTimestamp>,
	) {
		let start_recent_trade_timestamp = recent_trade_timestamp.cloned();
		let mut last_recent_trade_timestamp = recent_trade_timestamp.cloned();

		#[rustfmt::skip]
		let Some(
			realtime_databento,
		) = &mut self.realtime else {
			return;
		};

		loop {
			let next_record = realtime_databento
				.client
				.try_next_record();

			let record_ref = match next_record {
				Err(error) => {
					error!("{error}");
					break;
				}
				Ok(record_ref) => record_ref,
			};

			if let Some(record_ref) = record_ref {
				let Some(trade_trade_timestamp) = process_record_trades::<IS>(record_ref) else {
					continue;
				};

				last_recent_trade_timestamp = Some(*trade_trade_timestamp.trade_timestamp());
				instrument_data.new_trades([trade_trade_timestamp].into_iter());

				continue;
			}

			let size = match realtime_databento
				.client
				.fill_buf()
				.await
			{
				Err(error) => {
					error!("{error}");
					break;
				}
				Ok(size) => size,
			};

			if size == 0 {
				break;
			}

			if realtime_databento
				.client
				.is_closed()
			{
				unimplemented!();
			}
		}

		let Some(last_recent_trade_timestamp) = last_recent_trade_timestamp else {
			return;
		};

		let start_recent_trade_timestamp =
			start_recent_trade_timestamp.unwrap_or(last_recent_trade_timestamp);

		let tick_timestamp = TickTimestamp::new(*start_recent_trade_timestamp.timestamp());

		strategy
			.data_update(
				&tick_timestamp,
				instrument_data,
				order_manager,
				directional_exposure,
				active_order_goals,
				active_state_goal,
				deferred_order_actions,
				order_id_generator,
			)
			.await;
	}
}

impl<IS: InstrumentSpec + Send, OB: OrdersBackend<IS> + Send> DataBackend<IS>
	for Databento<IS, OB>
{
	fn new(key: Option<&str>) -> Self {
		Databento::new(key)
	}

	fn file_name_postpend() -> &'static str {
		".dbn"
	}

	async fn stream_in(
		&mut self,
		buf_reader: BufReader<File>,
	) -> impl ExactSizeIterator<Item = TradeTradeTimestamp<IS>> {
		let mut decoder = AsyncDbnDecoder::new(buf_reader)
			.await
			.unwrap();

		let mut trade_trade_timestamps: Vec<TradeTradeTimestamp<IS>> =
			Vec::with_capacity(1_000_000);

		while let Some(record_ref) = decoder
			.decode_record_ref()
			.await
			.unwrap()
		{
			let Some(trade_trade_timestamp) = process_record_trades(record_ref) else {
				continue;
			};

			trade_trade_timestamps.push(trade_trade_timestamp);
		}

		trade_trade_timestamps.into_iter()
	}
}

fn process_record_trades<IS: InstrumentSpec>(
	record_ref: RecordRef
) -> Option<TradeTradeTimestamp<IS>> {
	match record_ref.as_enum().unwrap() {
		RecordRefEnum::Mbo(mbo_msg) => unimplemented!(),
		RecordRefEnum::Mbp10(msg) => unimplemented!(),
		RecordRefEnum::Mbp1(msg) => unimplemented!(),
		RecordRefEnum::Bbo(msg) => unimplemented!(),
		RecordRefEnum::Cbbo(msg) => unimplemented!(),
		RecordRefEnum::Trade(trade_msg) => {
			let Ok(side) = trade_msg.side() else {
				return None;
			};

			let aggressor_side = match side {
				Side::Ask => AggressorSide::Ask,
				Side::Bid => AggressorSide::Bid,
				Side::None => return None,
			};

			let trade_trade_timestamp = create_trade_trade_timestamp::<IS>(
				trade_msg.ts_recv,
				trade_msg.price,
				trade_msg.size,
				aggressor_side,
			);

			Some(trade_trade_timestamp)
		}
		RecordRefEnum::Ohlcv(msg) => unimplemented!(),
		_ => None,
	}
}

fn create_trade_trade_timestamp<IS: InstrumentSpec>(
	ts_recv: u64,
	price: i64,
	volume: u32,
	aggressor_side: AggressorSide,
) -> TradeTradeTimestamp<IS> {
	let trade_timestamp = TradeTimestamp::new(Timestamp::new(
		ts_recv as i128,
	));

	let price = if IS::PriceType::IS_FLOATING_POINT {
		price_float_ticks::<IS>(price)
	} else {
		price_integer_ticks::<IS>(price)
	};

	let volume = AggressiveVolume::new(
		DirectionlessVolume::new_checked(IS::VolumeType::from_u64(
			volume as u64,
		))
		.unwrap(),
		aggressor_side,
	);

	let trade = Trade {
		price,
		aggressive_volume: volume,
	};

	TradeTradeTimestamp::new(trade, trade_timestamp)
}

#[inline]
fn price_integer_ticks<IS: InstrumentSpec>(price: i64) -> AbsolutePrice<IS> {
	let price_scale_ticks = FIXED_PRICE_SCALE / IS::PriceSpec::TICKS_PER_POINT as i64;
	let price = IS::PriceType::from_i64(price / price_scale_ticks);

	AbsolutePrice::new(price)
}

#[inline]
fn price_float_ticks<IS: InstrumentSpec>(price: i64) -> AbsolutePrice<IS> {
	let price_scale_ticks = FIXED_PRICE_SCALE as f64 / IS::PriceSpec::TICKS_PER_POINT as f64;
	let price = IS::PriceType::from_f64(price as f64 / price_scale_ticks);

	AbsolutePrice::new(price)
}