use chrono::{NaiveDate, NaiveDateTime};
use pine_builtin_macro::BuiltinFunction;
use pine_core::Bar;
use pine_interpreter::{
Builtin, BuiltinFn, EvaluatedArg, Interpreter, PineOutput, RuntimeError, Value,
};
use std::rc::Rc;
use std::time::{SystemTime, UNIX_EPOCH};
pub fn register_bar_time<O: PineOutput>(bar: &Bar) -> Value<O> {
Value::Number(bar.time as f64)
}
pub fn register_timenow<O: PineOutput>() -> Value<O> {
match SystemTime::now().duration_since(UNIX_EPOCH) {
Ok(since_epoch) => Value::Number(since_epoch.as_millis() as f64),
Err(_) => Value::Na,
}
}
fn ymd_to_millis(y: i64, mo: i64, d: i64, h: i64, mi: i64, s: i64) -> Option<i64> {
let date = NaiveDate::from_ymd_opt(y as i32, mo as u32, d as u32)?;
let dt = date.and_hms_opt(h as u32, mi as u32, s as u32)?;
Some(dt.and_utc().timestamp_millis())
}
fn parse_date_string(s: &str) -> Option<i64> {
let s = s.trim();
const DATETIME_FORMATS: &[&str] = &[
"%Y-%m-%dT%H:%M:%S",
"%Y-%m-%d %H:%M:%S",
"%Y-%m-%d %H:%M",
"%d %b %Y %H:%M:%S",
"%d %b %Y %H:%M",
];
for fmt in DATETIME_FORMATS {
if let Ok(dt) = NaiveDateTime::parse_from_str(s, fmt) {
return Some(dt.and_utc().timestamp_millis());
}
}
const DATE_FORMATS: &[&str] = &["%Y-%m-%d", "%d %b %Y", "%d %B %Y"];
for fmt in DATE_FORMATS {
if let Ok(date) = NaiveDate::parse_from_str(s, fmt) {
return Some(date.and_hms_opt(0, 0, 0)?.and_utc().timestamp_millis());
}
}
None
}
fn timestamp_fn<O: PineOutput>() -> BuiltinFn<O> {
Rc::new(|_ctx, call_args| {
let values: Vec<&Value<O>> = call_args
.args
.iter()
.map(|arg| match arg {
EvaluatedArg::Positional(v) => v,
EvaluatedArg::Named { value, .. } => value,
})
.collect();
if let [Value::String(s)] = values.as_slice() {
return Ok(parse_date_string(s)
.map(|ms| Value::Number(ms as f64))
.unwrap_or(Value::Na));
}
let nums: Vec<i64> = values
.iter()
.filter_map(|v| v.as_number().ok().map(|n| n as i64))
.collect();
if nums.len() < 3 {
return Ok(Value::Na);
}
let get = |i: usize| nums.get(i).copied().unwrap_or(0);
Ok(
ymd_to_millis(get(0), get(1), get(2), get(3), get(4), get(5))
.map(|ms| Value::Number(ms as f64))
.unwrap_or(Value::Na),
)
})
}
#[derive(BuiltinFunction)]
#[builtin(name = "year")]
struct Year {
time: f64,
}
impl Year {
fn execute<O: PineOutput>(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
let secs = (self.time / 1000.0) as i64;
let year = 1970 + (secs / (365 * 24 * 60 * 60));
Ok(Value::Number(year as f64))
}
}
#[derive(BuiltinFunction)]
#[builtin(name = "month")]
struct Month {
time: f64,
}
impl Month {
fn execute<O: PineOutput>(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
let secs = (self.time / 1000.0) as i64;
let days_since_epoch = secs / (24 * 60 * 60);
let month = ((days_since_epoch % 365) / 30) + 1;
let month = month.min(12);
Ok(Value::Number(month as f64))
}
}
#[derive(BuiltinFunction)]
#[builtin(name = "dayofmonth")]
struct DayOfMonth {
time: f64,
}
impl DayOfMonth {
fn execute<O: PineOutput>(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
let secs = (self.time / 1000.0) as i64;
let days_since_epoch = secs / (24 * 60 * 60);
let day = (days_since_epoch % 30) + 1;
Ok(Value::Number(day as f64))
}
}
#[derive(BuiltinFunction)]
#[builtin(name = "dayofweek")]
struct DayOfWeek {
time: f64,
}
impl DayOfWeek {
fn execute<O: PineOutput>(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
let secs = (self.time / 1000.0) as i64;
let days_since_epoch = secs / (24 * 60 * 60);
let day = ((days_since_epoch + 4) % 7) + 1;
Ok(Value::Number(day as f64))
}
}
#[derive(BuiltinFunction)]
#[builtin(name = "hour")]
struct Hour {
time: f64,
}
impl Hour {
fn execute<O: PineOutput>(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
let secs = (self.time / 1000.0) as i64;
let hour = (secs / 3600) % 24;
Ok(Value::Number(hour as f64))
}
}
#[derive(BuiltinFunction)]
#[builtin(name = "minute")]
struct Minute {
time: f64,
}
impl Minute {
fn execute<O: PineOutput>(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
let secs = (self.time / 1000.0) as i64;
let minute = (secs / 60) % 60;
Ok(Value::Number(minute as f64))
}
}
#[derive(BuiltinFunction)]
#[builtin(name = "second")]
struct Second {
time: f64,
}
impl Second {
fn execute<O: PineOutput>(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
let secs = (self.time / 1000.0) as i64;
let second = secs % 60;
Ok(Value::Number(second as f64))
}
}
pub fn register_time_functions<O: PineOutput>() -> Vec<(String, Value<O>)> {
vec![
(
"timestamp".to_string(),
Value::BuiltinFunction(Builtin::untyped(timestamp_fn::<O>())),
),
("year".to_string(), Year::builtin_value::<O>()),
("month".to_string(), Month::builtin_value::<O>()),
("dayofmonth".to_string(), DayOfMonth::builtin_value::<O>()),
("dayofweek".to_string(), DayOfWeek::builtin_value::<O>()),
("hour".to_string(), Hour::builtin_value::<O>()),
("minute".to_string(), Minute::builtin_value::<O>()),
("second".to_string(), Second::builtin_value::<O>()),
]
}