pine-builtins 0.2.2

use pine_builtin_macro::BuiltinFunction;
use pine_core::{
    AlertConditionOutput, BoxOutput, DrawingOutput, FillOutput, GlobalOutput, InputOutput,
    LabelOutput, LineOutput, LogOutput, MetadataOutput, PineOutput, PlotOutput, TableOutput,
};
use pine_core::{PineVersion, SymInfo, Timeframe};
use pine_interpreter::{Builtin, Interpreter, RuntimeError, Value};
use std::collections::HashMap;
use std::rc::Rc;

// Re-export for convenience
pub use pine_core::Bar;
pub use pine_core::DefaultPineOutput;
pub use pine_core::LogLevel;
pub use pine_interpreter::BuiltinFn;
pub use pine_interpreter::EvaluatedArg;

// Namespace modules
mod alertcondition;
mod array;
mod barstate;
mod r#box;
mod chart;
mod color;
mod constants;
mod currency;
mod dividends;
mod earnings;
mod fill;
mod footprint;
mod globals;
mod indicator;
mod input;
mod label;
mod library;
mod line;
mod linefill;
mod log;
mod map;
mod math;
mod matrix;
mod plot;
mod polyline;
mod request;
mod runtime;
mod session;
mod str;
mod strategy;
mod syminfo;
mod ta;
mod table;
mod ticker;
mod time;
mod timeframe;

// Global utility functions - defined first so they can be referenced in register function

/// na(value) - Returns true if the value is na, false otherwise
#[derive(BuiltinFunction)]
#[builtin(name = "na")]
struct Na<O: PineOutput> {
    value: Value<O>,
}

impl<O: PineOutput> Na<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        Ok(Value::Bool(matches!(self.value, Value::Na)))
    }
}

/// bool(x) - Converts value to bool
#[derive(BuiltinFunction)]
#[builtin(name = "bool")]
struct Bool<O: PineOutput> {
    x: Value<O>,
}

impl<O: PineOutput> Bool<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        match &self.x {
            Value::Bool(b) => Ok(Value::Bool(*b)),
            Value::Int(n) => Ok(Value::Bool(*n != 0)),
            Value::Number(n) => Ok(Value::Bool(*n != 0.0)),
            Value::Na => Ok(Value::Bool(false)),
            _ => Ok(Value::Bool(true)),
        }
    }
}

/// int(x) - Converts value to int (truncates float)
#[derive(BuiltinFunction)]
#[builtin(name = "int")]
struct Int<O: PineOutput> {
    x: Value<O>,
}

impl<O: PineOutput> Int<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        match &self.x {
            Value::Int(n) => Ok(Value::Int(*n)),
            Value::Number(n) => Ok(Value::Int(n.trunc() as i64)),
            Value::Bool(b) => Ok(Value::Int(if *b { 1 } else { 0 })),
            Value::Na => Ok(Value::Na),
            _ => Err(RuntimeError::TypeError(format!(
                "Cannot convert {:?} to int",
                self.x
            ))),
        }
    }
}

/// float(x) - Converts value to float
#[derive(BuiltinFunction)]
#[builtin(name = "float")]
struct Float<O: PineOutput> {
    x: Value<O>,
}

impl<O: PineOutput> Float<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        match &self.x {
            Value::Int(n) => Ok(Value::Number(*n as f64)),
            Value::Number(n) => Ok(Value::Number(*n)),
            Value::Bool(b) => Ok(Value::Number(if *b { 1.0 } else { 0.0 })),
            Value::Na => Ok(Value::Na),
            _ => Err(RuntimeError::TypeError(format!(
                "Cannot convert {:?} to float",
                self.x
            ))),
        }
    }
}

/// nz(source, replacement) - Replaces na values with default or replacement value
#[derive(BuiltinFunction)]
#[builtin(name = "nz")]
struct Nz<O: PineOutput> {
    source: Value<O>,
    #[arg(default = Value::Number(0.0))]
    replacement: Value<O>,
}

impl<O: PineOutput> Nz<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        // na source -> the replacement (any type; defaults to 0 when omitted).
        // Pine `na` reaches here as `Value::Na` or a NaN number.
        match &self.source {
            Value::Na => Ok(self.replacement.clone()),
            Value::Number(n) if n.is_nan() => Ok(self.replacement.clone()),
            _ => Ok(self.source.clone()),
        }
    }
}

/// fixnan(source) - Replaces NaN values with previous nearest non-NaN value
#[derive(BuiltinFunction)]
#[builtin(name = "fixnan")]
struct Fixnan<O: PineOutput> {
    source: Value<O>,
}

impl<O: PineOutput> Fixnan<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        // This is a simplified implementation
        // A full implementation would need to track previous values across bar evaluations
        match &self.source {
            Value::Na => {
                // Try to get the last non-na value from context
                // For now, just return 0.0 as a placeholder
                Ok(Value::Number(0.0))
            }
            Value::Number(n) if n.is_nan() => Ok(Value::Number(0.0)),
            _ => Ok(self.source.clone()),
        }
    }
}

/// The na-cast functions (`box(x)`, `color(x)`, `string(x)`, …) — they cast `na`
/// to a type, which in our dynamic typing is the identity on the argument.
fn na_cast<O: PineOutput>() -> BuiltinFn<O> {
    Rc::new(|_ctx, call_args| {
        Ok(match call_args.args.into_iter().next() {
            Some(EvaluatedArg::Positional(v)) => v,
            Some(EvaluatedArg::Named { value, .. }) => value,
            None => Value::Na,
        })
    })
}

/// Makes a namespace object also callable as its type's na-cast (`box.new` and
/// `box(x)` on the same name).
fn callable_namespace<O: PineOutput>(namespace: Value<O>) -> Value<O> {
    match namespace {
        Value::Object {
            type_name,
            fields,
            value,
            ..
        } => Value::Object {
            type_name,
            fields,
            value,
            call: Some(Builtin::untyped(na_cast::<O>())),
        },
        other => other,
    }
}

/// Register all builtin namespaces as objects and global functions
/// Returns namespace objects to be loaded as variables (e.g., "array", "str", "ta")
/// and global builtin functions (e.g., "na")
/// Each member stores the builtin function pointer as Value::BuiltinFunction
///
/// This uses DefaultPineOutput for now. Full generic support will be added when the
/// BuiltinFunction macro is updated to support generic output types.
pub fn register_namespace_objects<
    O: PineOutput
        + LogOutput
        + PlotOutput
        + LabelOutput
        + BoxOutput
        + InputOutput
        + LineOutput
        + TableOutput
        + MetadataOutput
        + GlobalOutput
        + AlertConditionOutput
        + FillOutput
        + DrawingOutput,
>(
    version: PineVersion,
    syminfo: Option<SymInfo>,
    timeframe: Option<Timeframe>,
) -> (
    HashMap<String, Value<O>>,
    Vec<pine_interpreter::PerBarAdvance<O>>,
) {
    let mut namespaces = HashMap::new();
    let mut advances = Vec::new();

    // `syminfo` and `timeframe` are always present in Pine, so an absent one
    // falls back to defaults.
    namespaces.insert(
        "syminfo".to_string(),
        syminfo::create_syminfo(syminfo.unwrap_or_default()),
    );
    namespaces.insert(
        "timeframe".to_string(),
        timeframe::register(timeframe.unwrap_or_default()),
    );

    // Register namespace objects
    namespaces.insert("array".to_string(), array::register());
    namespaces.insert("box".to_string(), callable_namespace(r#box::register()));
    namespaces.insert("chart".to_string(), chart::register());
    namespaces.insert("color".to_string(), callable_namespace(color::register()));
    namespaces.insert("map".to_string(), map::register());
    namespaces.insert("session".to_string(), session::register());
    namespaces.insert("runtime".to_string(), runtime::register());
    namespaces.insert("alert".to_string(), alertcondition::register_alert());
    namespaces.insert("ticker".to_string(), ticker::register());
    namespaces.insert("earnings".to_string(), earnings::register());
    namespaces.insert("footprint".to_string(), footprint::register_footprint());
    namespaces.insert("volume_row".to_string(), footprint::register_volume_row());
    namespaces.insert("dividends".to_string(), dividends::register());
    namespaces.insert("currency".to_string(), currency::register());
    for (name, value) in input::register(version) {
        namespaces.insert(name, value);
    }
    namespaces.insert("label".to_string(), callable_namespace(label::register()));
    for (name, value) in line::register(version) {
        // `line` is also the `line(x)` na-cast; `hline` stays as-is.
        let value = if name == "line" {
            callable_namespace(value)
        } else {
            value
        };
        namespaces.insert(name, value);
    }
    namespaces.insert(
        "string".to_string(),
        Value::BuiltinFunction(Builtin::untyped(na_cast::<O>())),
    );
    namespaces.insert(
        "linefill".to_string(),
        callable_namespace(linefill::register()),
    );
    namespaces.insert("polyline".to_string(), polyline::register());
    namespaces.insert("table".to_string(), callable_namespace(table::register()));
    for (name, value) in indicator::register(version) {
        namespaces.insert(name, value);
    }
    for (name, value) in library::register(version) {
        namespaces.insert(name, value);
    }
    namespaces.insert("request".to_string(), request::register());
    namespaces.insert("strategy".to_string(), strategy::register(version));
    namespaces.insert("alertcondition".to_string(), alertcondition::register());
    namespaces.insert("fill".to_string(), fill::register());
    for (name, value) in globals::register() {
        namespaces.insert(name, value);
    }

    // Constant-only namespaces (string tags used as arguments elsewhere).
    namespaces.insert("size".to_string(), constants::size::register());
    namespaces.insert("shape".to_string(), constants::shape::register());
    namespaces.insert("location".to_string(), constants::location::register());
    namespaces.insert("position".to_string(), constants::position::register());
    namespaces.insert("display".to_string(), constants::display::register());
    namespaces.insert("format".to_string(), constants::format::register());
    namespaces.insert("order".to_string(), constants::order::register());
    namespaces.insert("text".to_string(), constants::text::register());
    namespaces.insert("xloc".to_string(), constants::xloc::register());
    namespaces.insert("extend".to_string(), constants::extend::register());
    namespaces.insert("barmerge".to_string(), constants::barmerge::register());
    namespaces.insert("yloc".to_string(), constants::yloc::register());
    namespaces.insert("scale".to_string(), constants::scale::register());
    namespaces.insert("font".to_string(), constants::font::register());
    namespaces.insert("splits".to_string(), constants::splits::register());
    namespaces.insert("adjustment".to_string(), constants::adjustment::register());
    namespaces.insert(
        "backadjustment".to_string(),
        constants::backadjustment::register(),
    );
    namespaces.insert(
        "settlement_as_close".to_string(),
        constants::settlement_as_close::register(),
    );
    namespaces.insert("log".to_string(), log::register());
    for (name, func) in math::register(version) {
        namespaces.insert(name, func);
    }
    namespaces.insert("matrix".to_string(), matrix::register());
    for (name, func) in str::register(version) {
        namespaces.insert(name, func);
    }
    let (ta_ns, ta_advance) = ta::register(version);
    for (name, func) in ta_ns {
        namespaces.insert(name, func);
    }
    advances.push(ta_advance);

    // Register global builtin functions
    namespaces.insert("na".to_string(), Na::<O>::builtin_value());
    namespaces.insert("bool".to_string(), Bool::<O>::builtin_value());
    namespaces.insert("int".to_string(), Int::<O>::builtin_value());
    namespaces.insert("float".to_string(), Float::<O>::builtin_value());
    namespaces.insert("nz".to_string(), Nz::<O>::builtin_value());
    namespaces.insert("fixnan".to_string(), Fixnan::<O>::builtin_value());

    // Register time/date functions
    for (name, func) in time::register_time_functions() {
        namespaces.insert(name, func);
    }
    // `dayofweek` is value + function + namespace at once; its scalar is
    // refreshed each bar via `per_bar_object_values`.
    namespaces.insert("dayofweek".to_string(), time::register_dayofweek());
    namespaces.insert("time_close".to_string(), time::register_time_close());
    namespaces.insert(
        "time_tradingday".to_string(),
        time::register_time_tradingday(),
    );

    // Register plot functions
    for (name, func) in plot::register_plot_functions() {
        namespaces.insert(name, func);
    }

    (namespaces, advances)
}

/// Per-bar variables, rebuilt for each [`Bar`] and registered before it executes.
///
/// The compile-time counterpart is [`register_namespace_objects`]; this holds the
/// values that change every bar.
pub fn register_per_bar<O: PineOutput>(bar: &Bar) -> Vec<(String, Value<O>)> {
    vec![
        ("barstate".to_string(), barstate::register(bar)),
        ("timenow".to_string(), time::register_timenow()),
    ]
}

/// Every built-in variable a [`Bar`] sets: the price series (OHLCV and its
/// standard derivations) as history-carrying [`Value::Series`], `bar_index` as a
/// plain number, then the per-bar namespaces from [`register_per_bar`].
///
/// The single source of truth for these names and formulas, shared by execution
/// and the sema symbol table. The value's kind says how to store it: a
/// `Value::Series` is one the interpreter should advance to accumulate lookback
/// (`close[1]`); everything else is a plain assignment. Sema, needing only the
/// names to resolve, registers them as-is.
pub fn per_bar_variables<O: PineOutput>(
    bar: &Bar,
    last_bar: Option<&Bar>,
) -> Vec<(String, Value<O>)> {
    let series = |id: &str, value: f64| {
        (
            id.to_string(),
            Value::Series(pine_interpreter::Series {
                id: id.to_string(),
                current: Box::new(Value::Number(value)),
                history: None,
            }),
        )
    };
    let mut vars = vec![
        series("open", bar.open),
        series("high", bar.high),
        series("low", bar.low),
        series("close", bar.close),
        series("volume", bar.volume),
        series("hl2", (bar.high + bar.low) / 2.0),
        series("hlc3", (bar.high + bar.low + bar.close) / 3.0),
        series("hlcc4", (bar.high + bar.low + bar.close * 2.0) / 4.0),
        series("ohlc4", (bar.open + bar.high + bar.low + bar.close) / 4.0),
        ("bar_index".to_string(), Value::Number(bar.index as f64)),
        (
            "last_bar_index".to_string(),
            last_bar.map_or(Value::Na, |b| Value::Number(b.index as f64)),
        ),
        (
            "last_bar_time".to_string(),
            last_bar.map_or(Value::Na, |b| Value::Number(b.time as f64)),
        ),
    ];
    vars.extend(register_per_bar(bar));
    vars
}

#[cfg(test)]
mod tests {
    use super::*;
    use pine_interpreter::{EvaluatedArg, FunctionCallArgs};

    #[test]
    fn test_na() {
        let mut ctx = Interpreter::<DefaultPineOutput>::new();

        // Test with na value
        let args = vec![EvaluatedArg::Positional(Value::Na)];
        let result = Na::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Bool(true));

        // Test with number
        let args = vec![EvaluatedArg::Positional(Value::Number(42.0))];
        let result = Na::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Bool(false));

        // Test with string
        let args = vec![EvaluatedArg::Positional(Value::String("hello".to_string()))];
        let result = Na::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Bool(false));

        // Test with bool
        let args = vec![EvaluatedArg::Positional(Value::Bool(true))];
        let result = Na::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Bool(false));
    }

    #[test]
    fn test_bool() {
        let mut ctx = Interpreter::<DefaultPineOutput>::new();

        // Test number to bool
        let args = vec![EvaluatedArg::Positional(Value::Number(5.0))];
        let result = Bool::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Bool(true));

        let args = vec![EvaluatedArg::Positional(Value::Number(0.0))];
        let result = Bool::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Bool(false));

        // Test na to bool
        let args = vec![EvaluatedArg::Positional(Value::Na)];
        let result = Bool::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Bool(false));
    }

    #[test]
    fn test_int() {
        let mut ctx = Interpreter::<DefaultPineOutput>::new();

        // Test float to int (truncate)
        let args = vec![EvaluatedArg::Positional(Value::Number(5.7))];
        let result = Int::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(5.0));

        let args = vec![EvaluatedArg::Positional(Value::Number(-5.7))];
        let result = Int::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(-5.0));

        // Test bool to int
        let args = vec![EvaluatedArg::Positional(Value::Bool(true))];
        let result = Int::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(1.0));

        // Test na to int
        let args = vec![EvaluatedArg::Positional(Value::Na)];
        let result = Int::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Na);
    }

    #[test]
    fn test_float() {
        let mut ctx = Interpreter::<DefaultPineOutput>::new();

        // Test number to float
        let args = vec![EvaluatedArg::Positional(Value::Number(5.0))];
        let result = Float::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(5.0));

        // Test bool to float
        let args = vec![EvaluatedArg::Positional(Value::Bool(true))];
        let result = Float::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(1.0));

        // Test na to float
        let args = vec![EvaluatedArg::Positional(Value::Na)];
        let result = Float::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Na);
    }

    #[test]
    fn test_nz() {
        let mut ctx = Interpreter::<DefaultPineOutput>::new();

        // Test na value without replacement (should return 0.0)
        let args = vec![EvaluatedArg::Positional(Value::Na)];
        let result = Nz::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(0.0));

        // Test na value with replacement
        let args = vec![
            EvaluatedArg::Positional(Value::Na),
            EvaluatedArg::Positional(Value::Number(42.0)),
        ];
        let result = Nz::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(42.0));

        // Test non-na value (should return source)
        let args = vec![EvaluatedArg::Positional(Value::Number(5.0))];
        let result = Nz::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(5.0));
    }

    #[test]
    fn test_fixnan() {
        let mut ctx = Interpreter::<DefaultPineOutput>::new();

        // Test na value
        let args = vec![EvaluatedArg::Positional(Value::Na)];
        let result = Fixnan::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(0.0));

        // Test normal value
        let args = vec![EvaluatedArg::Positional(Value::Number(5.0))];
        let result = Fixnan::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(5.0));

        // Test NaN value
        let args = vec![EvaluatedArg::Positional(Value::Number(f64::NAN))];
        let result = Fixnan::builtin_fn(&mut ctx, FunctionCallArgs::without_types(args)).unwrap();
        assert_eq!(result, Value::Number(0.0));
    }
}