pine-builtins 0.1.0

use pine_builtin_macro::BuiltinFunction;
use pine_interpreter::{Interpreter, PineOutput, RuntimeError, Value};
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

/// matrix.new<type>() - Creates a new typed matrix
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.new", type_params = 1)]
struct MatrixNew<O: PineOutput> {
    #[type_param]
    element_type: String,
    #[arg(default = 0.0)]
    rows: f64,
    #[arg(default = 0.0)]
    columns: f64,
    #[arg(default = Value::Na)]
    initial_value: Value<O>,
}

impl<O: PineOutput> MatrixNew<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let rows = self.rows as usize;
        let columns = self.columns as usize;

        // Validate element type
        if !matches!(
            self.element_type.as_str(),
            "int" | "float" | "string" | "bool"
        ) {
            return Err(RuntimeError::TypeError(format!(
                "Invalid matrix element type '{}'. Must be int, float, string, or bool",
                self.element_type
            )));
        }

        // Create a matrix filled with the initial value
        let mut matrix_data = Vec::with_capacity(rows);
        for _ in 0..rows {
            let mut row = Vec::with_capacity(columns);
            for _ in 0..columns {
                row.push(self.initial_value.clone());
            }
            matrix_data.push(row);
        }

        Ok(Value::Matrix {
            element_type: self.element_type.clone(),
            data: Rc::new(RefCell::new(matrix_data)),
        })
    }
}

/// matrix.get() - Gets an element from the matrix
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.get")]
struct MatrixGet<O: PineOutput> {
    id: Value<O>,
    row: f64,
    column: f64,
}

impl<O: PineOutput> MatrixGet<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let matrix = match &self.id {
            Value::Matrix { data, .. } => data,
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let row_idx = self.row as usize;
        let col_idx = self.column as usize;

        let matrix_ref = matrix.borrow();
        if row_idx >= matrix_ref.len() {
            return Err(RuntimeError::IndexOutOfBounds(row_idx));
        }
        if col_idx >= matrix_ref[row_idx].len() {
            return Err(RuntimeError::IndexOutOfBounds(col_idx));
        }

        Ok(matrix_ref[row_idx][col_idx].clone())
    }
}

/// matrix.set() - Sets an element in the matrix
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.set")]
struct MatrixSet<O: PineOutput> {
    id: Value<O>,
    row: f64,
    column: f64,
    value: Value<O>,
}

impl<O: PineOutput> MatrixSet<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let matrix = match &self.id {
            Value::Matrix { data, .. } => data,
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let row_idx = self.row as usize;
        let col_idx = self.column as usize;

        let mut matrix_ref = matrix.borrow_mut();
        if row_idx >= matrix_ref.len() {
            return Err(RuntimeError::IndexOutOfBounds(row_idx));
        }
        if col_idx >= matrix_ref[row_idx].len() {
            return Err(RuntimeError::IndexOutOfBounds(col_idx));
        }

        matrix_ref[row_idx][col_idx] = self.value.clone();
        Ok(Value::Na)
    }
}

/// matrix.rows() - Returns the number of rows
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.rows")]
struct MatrixRows<O: PineOutput> {
    id: Value<O>,
}

impl<O: PineOutput> MatrixRows<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let matrix = match &self.id {
            Value::Matrix { data, .. } => data,
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let count = matrix.borrow().len();
        Ok(Value::Number(count as f64))
    }
}

/// matrix.columns() - Returns the number of columns
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.columns")]
struct MatrixColumns<O: PineOutput> {
    id: Value<O>,
}

impl<O: PineOutput> MatrixColumns<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let matrix = match &self.id {
            Value::Matrix { data, .. } => data,
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let matrix_ref = matrix.borrow();
        let count = if matrix_ref.is_empty() {
            0
        } else {
            matrix_ref[0].len()
        };
        Ok(Value::Number(count as f64))
    }
}

/// matrix.elements_count() - Returns total number of elements
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.elements_count")]
struct MatrixElementsCount<O: PineOutput> {
    id: Value<O>,
}

impl<O: PineOutput> MatrixElementsCount<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let matrix = match &self.id {
            Value::Matrix { data, .. } => data,
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let matrix_ref = matrix.borrow();
        let count: usize = matrix_ref.iter().map(|row| row.len()).sum();
        Ok(Value::Number(count as f64))
    }
}

/// matrix.fill() - Fills the matrix with a value
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.fill")]
struct MatrixFill<O: PineOutput> {
    id: Value<O>,
    value: Value<O>,
}

impl<O: PineOutput> MatrixFill<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let matrix = match &self.id {
            Value::Matrix { data, .. } => data,
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let mut matrix_ref = matrix.borrow_mut();
        for row in matrix_ref.iter_mut() {
            for cell in row.iter_mut() {
                *cell = self.value.clone();
            }
        }

        Ok(Value::Na)
    }
}

/// matrix.copy() - Creates a copy of the matrix
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.copy")]
struct MatrixCopy<O: PineOutput> {
    id: Value<O>,
}

impl<O: PineOutput> MatrixCopy<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let (matrix, element_type) = match &self.id {
            Value::Matrix { data, element_type } => (data, element_type.clone()),
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let matrix_ref = matrix.borrow();
        let copied_data = matrix_ref.clone();
        Ok(Value::Matrix {
            element_type,
            data: Rc::new(RefCell::new(copied_data)),
        })
    }
}

/// matrix.add_row() - Adds a row to the matrix
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.add_row")]
struct MatrixAddRow<O: PineOutput> {
    id: Value<O>,
    row: f64,
    #[arg(default = Value::Na)]
    array_id: Value<O>,
}

impl<O: PineOutput> MatrixAddRow<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let matrix = match &self.id {
            Value::Matrix { data, .. } => data,
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let row_idx = self.row as usize;

        let mut matrix_ref = matrix.borrow_mut();
        let cols = if matrix_ref.is_empty() {
            0
        } else {
            matrix_ref[0].len()
        };

        // If array_id is provided, use its values, otherwise use na
        let new_row = match &self.array_id {
            Value::Array(arr) => arr.borrow().clone(),
            Value::Na => vec![Value::Na; cols],
            _ => {
                return Err(RuntimeError::TypeError(
                    "array_id must be an array".to_string(),
                ))
            }
        };

        if row_idx > matrix_ref.len() {
            return Err(RuntimeError::IndexOutOfBounds(row_idx));
        }

        matrix_ref.insert(row_idx, new_row);
        Ok(Value::Na)
    }
}

/// matrix.add_col() - Adds a column to the matrix
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.add_col")]
struct MatrixAddCol<O: PineOutput> {
    id: Value<O>,
    column: f64,
    #[arg(default = Value::Na)]
    array_id: Value<O>,
}

impl<O: PineOutput> MatrixAddCol<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let matrix = match &self.id {
            Value::Matrix { data, .. } => data,
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let col_idx = self.column as usize;

        let mut matrix_ref = matrix.borrow_mut();

        // Get column values from array or use na
        let col_values = match &self.array_id {
            Value::Array(arr) => arr.borrow().clone(),
            Value::Na => vec![Value::Na; matrix_ref.len()],
            _ => {
                return Err(RuntimeError::TypeError(
                    "array_id must be an array".to_string(),
                ))
            }
        };

        for (i, row) in matrix_ref.iter_mut().enumerate() {
            if col_idx > row.len() {
                return Err(RuntimeError::IndexOutOfBounds(col_idx));
            }
            let val = col_values.get(i).cloned().unwrap_or(Value::Na);
            row.insert(col_idx, val);
        }

        Ok(Value::Na)
    }
}

/// matrix.transpose() - Transposes the matrix
#[derive(BuiltinFunction)]
#[builtin(name = "matrix.transpose")]
struct MatrixTranspose<O: PineOutput> {
    id: Value<O>,
}

impl<O: PineOutput> MatrixTranspose<O> {
    fn execute(&self, _ctx: &mut Interpreter<O>) -> Result<Value<O>, RuntimeError> {
        let (matrix, element_type) = match &self.id {
            Value::Matrix { data, element_type } => (data, element_type.clone()),
            _ => return Err(RuntimeError::TypeError("Expected matrix".to_string())),
        };

        let matrix_ref = matrix.borrow();
        if matrix_ref.is_empty() {
            return Ok(Value::Matrix {
                element_type,
                data: Rc::new(RefCell::new(vec![])),
            });
        }

        let rows = matrix_ref.len();
        let cols = matrix_ref[0].len();
        let mut transposed = vec![vec![Value::Na; rows]; cols];

        for i in 0..rows {
            for j in 0..cols {
                transposed[j][i] = matrix_ref[i][j].clone();
            }
        }

        Ok(Value::Matrix {
            element_type,
            data: Rc::new(RefCell::new(transposed)),
        })
    }
}

/// Register the matrix namespace with all functions
pub fn register<O: PineOutput>() -> Value<O> {
    let mut members: HashMap<String, Value<O>> = HashMap::new();

    members.insert("new".to_string(), MatrixNew::<O>::builtin_value());
    members.insert("get".to_string(), MatrixGet::<O>::builtin_value());
    members.insert("set".to_string(), MatrixSet::<O>::builtin_value());
    members.insert("rows".to_string(), MatrixRows::<O>::builtin_value());
    members.insert("columns".to_string(), MatrixColumns::<O>::builtin_value());
    members.insert(
        "elements_count".to_string(),
        MatrixElementsCount::<O>::builtin_value(),
    );
    members.insert("fill".to_string(), MatrixFill::<O>::builtin_value());
    members.insert("copy".to_string(), MatrixCopy::<O>::builtin_value());
    members.insert("add_row".to_string(), MatrixAddRow::<O>::builtin_value());
    members.insert("add_col".to_string(), MatrixAddCol::<O>::builtin_value());
    members.insert(
        "transpose".to_string(),
        MatrixTranspose::<O>::builtin_value(),
    );

    Value::Object {
        type_name: "matrix".to_string(),
        fields: Rc::new(RefCell::new(members)),
        call: None,
    }
}