import pl, { DataFrame, Series } from 'nodejs-polars'
import fs from 'fs'

export interface CSVDatabase {
  df: DataFrame
  filepath: string
}

export function create_csv_database(filepath: string, initialData: any): CSVDatabase {
  let df = DataFrame(initialData)
  df.writeCSV(filepath)
  return { df, filepath }
}

export function save_db(db: CSVDatabase) {
  db.df.writeCSV(db.filepath)
}

export function add_row(db: CSVDatabase, row: any) {
  let d = DataFrame([row])
  db.df = db.df.vstack(d)
}

export function print_head(db: CSVDatabase) {
  console.log(db.df.head())
}

export function delete_row_by_key(db: CSVDatabase, key: string, keyColumn: string) {
  db.df = db.df.filter(pl.col(keyColumn).neq(pl.lit(key)))
}

export function select_columns(db: CSVDatabase, columns: string[]): DataFrame {
  return db.df.select(...columns)
}

export function join_dataframes(db1: CSVDatabase, db2: CSVDatabase, on: string): DataFrame {
  return db1.df.join(db2.df, { on, how: "inner" })
}

export function add_simple_index(db: CSVDatabase) {
  const indices = Array.from(Array(db.df.height).keys()) // Create 0-N index
  db.df = db.df.withColumn(pl.Series("index", indices))
}