#include "quickjs.h"
#include <string.h>
#include <math.h>

#ifdef __APPLE__
#include <Accelerate/Accelerate.h>
#else
#include <cblas.h>
#include <lapacke.h>
#endif

static JSClassID js_matrix_class_id;
static JSClassID js_array_class_id;

typedef struct {
    double *data;
    int rows;
    int cols;
} matrix_t;

typedef struct {
    double *data;
    int size;
} array_t;

static void js_matrix_finalizer(JSRuntime *rt, JSValue val) {
    matrix_t *mat = JS_GetOpaque(val, js_matrix_class_id);
    if (mat) {
        js_free_rt(rt, mat->data);
        js_free_rt(rt, mat);
    }
}

static void js_array_finalizer(JSRuntime *rt, JSValue val) {
    array_t *arr = JS_GetOpaque(val, js_array_class_id);
    if (arr) {
        js_free_rt(rt, arr->data);
        js_free_rt(rt, arr);
    }
}

static JSClassDef js_matrix_class = {
    "Matrix",
    .finalizer = js_matrix_finalizer,
};

static JSClassDef js_array_class = {
    "Array",
    .finalizer = js_array_finalizer,
};

static matrix_t *js2matrix(JSContext *ctx, JSValue v) {
    return JS_GetOpaque(v, js_matrix_class_id);
}

static array_t *js2array(JSContext *ctx, JSValue v) {
    return JS_GetOpaque(v, js_array_class_id);
}

static JSValue js_matrix_constructor(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "Matrix constructor requires an array argument");
    
    if (!JS_IsArray(ctx, argv[0]))
        return JS_ThrowTypeError(ctx, "Matrix constructor requires an array argument");
    
    JSValue length_val = JS_GetPropertyStr(ctx, argv[0], "length");
    int32_t rows;
    JS_ToInt32(ctx, &rows, length_val);
    JS_FreeValue(ctx, length_val);
    
    if (rows == 0)
        return JS_ThrowRangeError(ctx, "Matrix cannot have 0 rows");
    
    // Get first row to determine columns
    JSValue first_row = JS_GetPropertyUint32(ctx, argv[0], 0);
    if (!JS_IsArray(ctx, first_row)) {
        JS_FreeValue(ctx, first_row);
        return JS_ThrowTypeError(ctx, "Matrix rows must be arrays");
    }
    
    JSValue col_length_val = JS_GetPropertyStr(ctx, first_row, "length");
    int32_t cols;
    JS_ToInt32(ctx, &cols, col_length_val);
    JS_FreeValue(ctx, col_length_val);
    JS_FreeValue(ctx, first_row);
    
    if (cols == 0)
        return JS_ThrowRangeError(ctx, "Matrix cannot have 0 columns");
    
    matrix_t *mat = js_malloc(ctx, sizeof(matrix_t));
    if (!mat)
        return JS_EXCEPTION;
    
    mat->rows = rows;
    mat->cols = cols;
    mat->data = js_malloc(ctx, sizeof(double) * rows * cols);
    if (!mat->data) {
        js_free(ctx, mat);
        return JS_EXCEPTION;
    }
    
    // Fill matrix data
    for (int i = 0; i < rows; i++) {
        JSValue row = JS_GetPropertyUint32(ctx, argv[0], i);
        if (!JS_IsArray(ctx, row)) {
            js_free(ctx, mat->data);
            js_free(ctx, mat);
            JS_FreeValue(ctx, row);
            return JS_ThrowTypeError(ctx, "All matrix rows must be arrays");
        }
        
        JSValue row_length_val = JS_GetPropertyStr(ctx, row, "length");
        int32_t row_cols;
        JS_ToInt32(ctx, &row_cols, row_length_val);
        JS_FreeValue(ctx, row_length_val);
        
        if (row_cols != cols) {
            js_free(ctx, mat->data);
            js_free(ctx, mat);
            JS_FreeValue(ctx, row);
            return JS_ThrowTypeError(ctx, "All matrix rows must have the same length");
        }
        
        for (int j = 0; j < cols; j++) {
            JSValue elem = JS_GetPropertyUint32(ctx, row, j);
            double val;
            if (JS_ToFloat64(ctx, &val, elem)) {
                js_free(ctx, mat->data);
                js_free(ctx, mat);
                JS_FreeValue(ctx, elem);
                JS_FreeValue(ctx, row);
                return JS_EXCEPTION;
            }
            mat->data[i * cols + j] = val;
            JS_FreeValue(ctx, elem);
        }
        JS_FreeValue(ctx, row);
    }
    
    JSValue obj = JS_NewObjectClass(ctx, js_matrix_class_id);
    JS_SetOpaque(obj, mat);
    return obj;
}

static JSValue js_array_constructor(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "Array constructor requires an array argument");
    
    if (!JS_IsArray(ctx, argv[0]))
        return JS_ThrowTypeError(ctx, "Array constructor requires an array argument");
    
    JSValue length_val = JS_GetPropertyStr(ctx, argv[0], "length");
    int32_t size;
    JS_ToInt32(ctx, &size, length_val);
    JS_FreeValue(ctx, length_val);
    
    if (size == 0)
        return JS_ThrowRangeError(ctx, "Array cannot have 0 elements");
    
    array_t *arr = js_malloc(ctx, sizeof(array_t));
    if (!arr)
        return JS_EXCEPTION;
    
    arr->size = size;
    arr->data = js_malloc(ctx, sizeof(double) * size);
    if (!arr->data) {
        js_free(ctx, arr);
        return JS_EXCEPTION;
    }
    
    // Fill array data
    for (int i = 0; i < size; i++) {
        JSValue elem = JS_GetPropertyUint32(ctx, argv[0], i);
        double val;
        if (JS_ToFloat64(ctx, &val, elem)) {
            js_free(ctx, arr->data);
            js_free(ctx, arr);
            JS_FreeValue(ctx, elem);
            return JS_EXCEPTION;
        }
        arr->data[i] = val;
        JS_FreeValue(ctx, elem);
    }
    
    JSValue obj = JS_NewObjectClass(ctx, js_array_class_id);
    JS_SetOpaque(obj, arr);
    return obj;
}

static JSValue js_matrix_inverse(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    matrix_t *mat = js2matrix(ctx, this_val);
    if (!mat)
        return JS_EXCEPTION;
    
    if (mat->rows != mat->cols)
        return JS_ThrowTypeError(ctx, "Matrix must be square for inversion");
    
    int n = mat->rows;
    
    // Create a copy of the matrix
    matrix_t *result = js_malloc(ctx, sizeof(matrix_t));
    if (!result)
        return JS_EXCEPTION;
    
    result->rows = n;
    result->cols = n;
    result->data = js_malloc(ctx, sizeof(double) * n * n);
    if (!result->data) {
        js_free(ctx, result);
        return JS_EXCEPTION;
    }
    
    memcpy(result->data, mat->data, sizeof(double) * n * n);
    
    // Allocate pivot array
    int *ipiv = js_malloc(ctx, sizeof(int) * n);
    if (!ipiv) {
        js_free(ctx, result->data);
        js_free(ctx, result);
        return JS_EXCEPTION;
    }
    
    // LU decomposition
#ifdef __APPLE__
    // For macOS with ILP64, use 64-bit integers
    __LAPACK_int nn = n;
    __LAPACK_int info = 0;
    __LAPACK_int *ipiv_lapack = js_malloc(ctx, sizeof(__LAPACK_int) * n);
    if (!ipiv_lapack) {
        js_free(ctx, ipiv);
        js_free(ctx, result->data);
        js_free(ctx, result);
        return JS_EXCEPTION;
    }
    
    // LAPACK uses column-major, but we'll transpose the result
    __LAPACK_int lda = n;
    dgetrf_(&nn, &nn, result->data, &lda, ipiv_lapack, &info);
    
    if (info != 0) {
        js_free(ctx, ipiv_lapack);
        js_free(ctx, ipiv);
        js_free(ctx, result->data);
        js_free(ctx, result);
        return JS_ThrowInternalError(ctx, "Matrix LU decomposition failed");
    }
    
    // Compute inverse
    __LAPACK_int lwork = n * n;
    double *work = js_malloc(ctx, sizeof(double) * lwork);
    if (!work) {
        js_free(ctx, ipiv_lapack);
        js_free(ctx, ipiv);
        js_free(ctx, result->data);
        js_free(ctx, result);
        return JS_EXCEPTION;
    }
    
    dgetri_(&nn, result->data, &lda, ipiv_lapack, work, &lwork, &info);
    js_free(ctx, work);
    js_free(ctx, ipiv_lapack);
#else
    int info = LAPACKE_dgetrf(LAPACK_ROW_MAJOR, n, n, result->data, n, ipiv);
    if (info != 0) {
        js_free(ctx, ipiv);
        js_free(ctx, result->data);
        js_free(ctx, result);
        return JS_ThrowInternalError(ctx, "Matrix LU decomposition failed");
    }
    
    // Compute inverse
    info = LAPACKE_dgetri(LAPACK_ROW_MAJOR, n, result->data, n, ipiv);
#endif
    js_free(ctx, ipiv);
    
    if (info != 0) {
        js_free(ctx, result->data);
        js_free(ctx, result);
        return JS_ThrowInternalError(ctx, "Matrix inversion failed");
    }
    
    JSValue obj = JS_NewObjectClass(ctx, js_matrix_class_id);
    JS_SetOpaque(obj, result);
    return obj;
}

static JSValue js_matrix_multiply(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "multiply requires an argument");
    
    matrix_t *A = js2matrix(ctx, this_val);
    if (!A)
        return JS_EXCEPTION;
    
    // Check if multiplying by another matrix
    matrix_t *B = js2matrix(ctx, argv[0]);
    if (B) {
        if (A->cols != B->rows)
            return JS_ThrowTypeError(ctx, "Matrix dimensions incompatible for multiplication");
        
        matrix_t *C = js_malloc(ctx, sizeof(matrix_t));
        if (!C)
            return JS_EXCEPTION;
        
        C->rows = A->rows;
        C->cols = B->cols;
        C->data = js_malloc(ctx, sizeof(double) * C->rows * C->cols);
        if (!C->data) {
            js_free(ctx, C);
            return JS_EXCEPTION;
        }
        
        // C = A * B
        cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                    A->rows, B->cols, A->cols,
                    1.0, A->data, A->cols,
                    B->data, B->cols,
                    0.0, C->data, C->cols);
        
        JSValue obj = JS_NewObjectClass(ctx, js_matrix_class_id);
        JS_SetOpaque(obj, C);
        return obj;
    }
    
    // Check if multiplying by an array (matrix-vector multiplication)
    array_t *x = js2array(ctx, argv[0]);
    if (x) {
        if (A->cols != x->size)
            return JS_ThrowTypeError(ctx, "Matrix columns must match array size");
        
        array_t *y = js_malloc(ctx, sizeof(array_t));
        if (!y)
            return JS_EXCEPTION;
        
        y->size = A->rows;
        y->data = js_malloc(ctx, sizeof(double) * y->size);
        if (!y->data) {
            js_free(ctx, y);
            return JS_EXCEPTION;
        }
        
        // y = A * x
        cblas_dgemv(CblasRowMajor, CblasNoTrans,
                    A->rows, A->cols,
                    1.0, A->data, A->cols,
                    x->data, 1,
                    0.0, y->data, 1);
        
        JSValue obj = JS_NewObjectClass(ctx, js_array_class_id);
        JS_SetOpaque(obj, y);
        return obj;
    }
    
    return JS_ThrowTypeError(ctx, "multiply requires a Matrix or Array argument");
}

static JSValue js_matrix_to_array(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    matrix_t *mat = js2matrix(ctx, this_val);
    if (!mat)
        return JS_EXCEPTION;
    
    JSValue arr = JS_NewArray(ctx);
    if (JS_IsException(arr))
        return arr;
    
    for (int i = 0; i < mat->rows; i++) {
        JSValue row = JS_NewArray(ctx);
        if (JS_IsException(row)) {
            JS_FreeValue(ctx, arr);
            return row;
        }
        for (int j = 0; j < mat->cols; j++) {
            JS_SetPropertyUint32(ctx, row, j, JS_NewFloat64(ctx, mat->data[i * mat->cols + j]));
        }
        JS_SetPropertyUint32(ctx, arr, i, row);
    }
    return arr;
}

static JSValue js_array_to_array(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    array_t *arr = js2array(ctx, this_val);
    if (!arr)
        return JS_EXCEPTION;
    
    JSValue jsarr = JS_NewArray(ctx);
    for (int i = 0; i < arr->size; i++) {
        JS_SetPropertyUint32(ctx, jsarr, i, JS_NewFloat64(ctx, arr->data[i]));
    }
    return jsarr;
}

static JSValue js_array_dot(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "dot requires an argument");
    
    array_t *a = js2array(ctx, this_val);
    if (!a)
        return JS_EXCEPTION;
    
    array_t *b = js2array(ctx, argv[0]);
    if (!b)
        return JS_ThrowTypeError(ctx, "dot requires an Array argument");
    
    if (a->size != b->size)
        return JS_ThrowTypeError(ctx, "Arrays must have the same size for dot product");
    
    double result = cblas_ddot(a->size, a->data, 1, b->data, 1);
    return JS_NewFloat64(ctx, result);
}

static JSValue js_array_norm(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    array_t *arr = js2array(ctx, this_val);
    if (!arr)
        return JS_EXCEPTION;
    
    double norm = cblas_dnrm2(arr->size, arr->data, 1);
    return JS_NewFloat64(ctx, norm);
}

static JSValue js_array_multiply(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "multiply requires an argument");
    
    array_t *arr = js2array(ctx, this_val);
    if (!arr)
        return JS_EXCEPTION;
    
    // Create result array
    array_t *result = js_malloc(ctx, sizeof(array_t));
    if (!result)
        return JS_EXCEPTION;
    
    result->size = arr->size;
    result->data = js_malloc(ctx, sizeof(double) * result->size);
    if (!result->data) {
        js_free(ctx, result);
        return JS_EXCEPTION;
    }
    
    // Copy original data
    memcpy(result->data, arr->data, sizeof(double) * arr->size);
    
    // Check if multiplying by scalar
    if (JS_IsNumber(argv[0])) {
        double scalar;
        if (JS_ToFloat64(ctx, &scalar, argv[0]))
            return JS_EXCEPTION;
        
        // Multiply each element by scalar
        if (result->size <= 4)
          for (int i = 0; i < result->size; i++)
            result->data[i] *= scalar;
        else
          cblas_dscal(result->size, scalar, result->data, 1);
    } else {
        js_free(ctx, result->data);
        js_free(ctx, result);
        return JS_ThrowTypeError(ctx, "multiply requires a number argument");
    }
    
    JSValue obj = JS_NewObjectClass(ctx, js_array_class_id);
    JS_SetOpaque(obj, result);
    return obj;
}

static JSValue js_array_add(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "add requires an argument");
    
    array_t *arr = js2array(ctx, this_val);
    if (!arr)
        return JS_EXCEPTION;
    
    // Create result array
    array_t *result = js_malloc(ctx, sizeof(array_t));
    if (!result)
        return JS_EXCEPTION;
    
    result->size = arr->size;
    result->data = js_malloc(ctx, sizeof(double) * result->size);
    if (!result->data) {
        js_free(ctx, result);
        return JS_EXCEPTION;
    }
    
    // Copy original data
    memcpy(result->data, arr->data, sizeof(double) * arr->size);
    
    // Check if adding scalar
    if (JS_IsNumber(argv[0])) {
        double scalar;
        if (JS_ToFloat64(ctx, &scalar, argv[0]))
            return JS_EXCEPTION;
        
        // Add scalar to each element
        for (int i = 0; i < result->size; i++)
            result->data[i] += scalar;
    } else {
        // Check if adding another array
        array_t *other = js2array(ctx, argv[0]);
        if (other) {
            if (arr->size != other->size) {
                js_free(ctx, result->data);
                js_free(ctx, result);
                return JS_ThrowTypeError(ctx, "Arrays must have the same size for element-wise addition");
            }
            
            if (arr->size <= 4) {
              for (int i = 0; i < arr->size; i++)
                result->data[i] += other->data[i];
            } else
              cblas_daxpy(result->size, 1.0, other->data, 1, result->data, 1);
        } else {
            js_free(ctx, result->data);
            js_free(ctx, result);
            return JS_ThrowTypeError(ctx, "add requires a number or Array argument");
        }
    }
    
    JSValue obj = JS_NewObjectClass(ctx, js_array_class_id);
    JS_SetOpaque(obj, result);
    return obj;
}

static JSValue js_array_slice(JSContext *js, JSValueConst self, int argc, JSValueConst *argv)
{
  array_t *arr = js2array(js, self);
  if (!arr)
    return JS_EXCEPTION;
  
  int32_t start = 0;
  int32_t end = arr->size;
  
  // Parse start argument
  if (argc >= 1 && !JS_IsNull(argv[0])) {
    if (JS_ToInt32(js, &start, argv[0]))
      return JS_EXCEPTION;
    
    // Handle negative indices
    if (start < 0) {
      start = arr->size + start;
      if (start < 0) start = 0;
    } else if (start > arr->size) {
      start = arr->size;
    }
  }
  
  // Parse end argument
  if (argc >= 2 && !JS_IsNull(argv[1])) {
    if (JS_ToInt32(js, &end, argv[1]))
      return JS_EXCEPTION;
    
    // Handle negative indices
    if (end < 0) {
      end = arr->size + end;
      if (end < 0) end = 0;
    } else if (end > arr->size) {
      end = arr->size;
    }
  }
  
  // Ensure start <= end
  if (start > end) {
    end = start;  // Return empty array
  }
  
  // Create JavaScript array
  JSValue jsarr = JS_NewArray(js);
  if (JS_IsException(jsarr))
    return jsarr;
  
  // Fill with sliced values
  for (int32_t i = start; i < end; i++) {
    JS_SetPropertyUint32(js, jsarr, i - start, JS_NewFloat64(js, arr->data[i]));
  }
  
  return jsarr;
}

static JSValue js_array_multiplyf(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "multiplyf requires an argument");
    
    array_t *arr = js2array(ctx, this_val);
    if (!arr)
        return JS_EXCEPTION;
    
    // Check if multiplying by scalar
    if (JS_IsNumber(argv[0])) {
        double scalar;
        if (JS_ToFloat64(ctx, &scalar, argv[0]))
            return JS_EXCEPTION;
        
        // Multiply each element by scalar in-place
        if (arr->size <= 4) {
            for (int i = 0; i < arr->size; i++)
                arr->data[i] *= scalar;
        } else {
            cblas_dscal(arr->size, scalar, arr->data, 1);
        }
        
        return JS_DupValue(ctx, this_val);
    } else {
        return JS_ThrowTypeError(ctx, "multiplyf requires a number argument");
    }
}

static JSValue js_array_addf(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "addf requires an argument");
    
    array_t *arr = js2array(ctx, this_val);
    if (!arr)
        return JS_EXCEPTION;
    
    // Check if adding scalar
    if (JS_IsNumber(argv[0])) {
        double scalar;
        if (JS_ToFloat64(ctx, &scalar, argv[0]))
            return JS_EXCEPTION;
        
        // Add scalar to each element in-place
        for (int i = 0; i < arr->size; i++)
            arr->data[i] += scalar;
    } else {
        // Check if adding another array
        array_t *other = js2array(ctx, argv[0]);
        if (other) {
            if (arr->size != other->size) {
                return JS_ThrowTypeError(ctx, "Arrays must have the same size for element-wise addition");
            }
            
            // Add arrays element-wise in-place
            if (arr->size <= 4) {
                for (int i = 0; i < arr->size; i++)
                    arr->data[i] += other->data[i];
            } else {
                cblas_daxpy(arr->size, 1.0, other->data, 1, arr->data, 1);
            }
        } else {
            return JS_ThrowTypeError(ctx, "addf requires a number or Array argument");
        }
    }
    
    return JS_DupValue(ctx, this_val);
}



static const JSCFunctionListEntry js_matrix_proto_funcs[] = {
    JS_CFUNC_DEF("inverse", 0, js_matrix_inverse),
    JS_CFUNC_DEF("multiply", 1, js_matrix_multiply),
    JS_CFUNC_DEF("toArray", 0, js_matrix_to_array),
};

static const JSCFunctionListEntry js_array_proto_funcs[] = {
    JS_CFUNC_DEF("dot", 1, js_array_dot),
    JS_CFUNC_DEF("norm", 0, js_array_norm),
    JS_CFUNC_DEF("multiply", 1, js_array_multiply),
    JS_CFUNC_DEF("add", 1, js_array_add),
    JS_CFUNC_DEF("multiplyf", 1, js_array_multiplyf),
    JS_CFUNC_DEF("addf", 1, js_array_addf),
    JS_CFUNC_DEF("toArray", 0, js_array_to_array),
    JS_CFUNC_DEF("toJSON",  0, js_array_to_array),  /* ← for JSON.stringify */
    JS_CFUNC_DEF("slice", 2, js_array_slice),
};

JSValue js_num_use(JSContext *ctx) {
    // Register Matrix class
    JS_NewClassID(&js_matrix_class_id);
    JS_NewClass(JS_GetRuntime(ctx), js_matrix_class_id, &js_matrix_class);
    JSValue matrix_proto = JS_NewObject(ctx);
    JS_SetPropertyFunctionList(ctx, matrix_proto, js_matrix_proto_funcs, 
                               sizeof(js_matrix_proto_funcs) / sizeof(JSCFunctionListEntry));
    JS_SetClassProto(ctx, js_matrix_class_id, matrix_proto);
    
    // Create Matrix constructor
    JSValue matrix_ctor = JS_NewCFunction2(ctx, js_matrix_constructor, "Matrix", 1, JS_CFUNC_constructor, 0);
    JS_SetConstructor(ctx, matrix_ctor, matrix_proto);
    
    // Register Array class
    JS_NewClassID(&js_array_class_id);
    JS_NewClass(JS_GetRuntime(ctx), js_array_class_id, &js_array_class);
    JSValue array_proto = JS_NewObject(ctx);
    JS_SetPropertyFunctionList(ctx, array_proto, js_array_proto_funcs,
                               sizeof(js_array_proto_funcs) / sizeof(JSCFunctionListEntry));
    JS_SetClassProto(ctx, js_array_class_id, array_proto);
    
    // Create Array constructor
    JSValue array_ctor = JS_NewCFunction2(ctx, js_array_constructor, "Array", 1, JS_CFUNC_constructor, 0);
    JS_SetConstructor(ctx, array_ctor, array_proto);
    
    JSValue export = JS_NewObject(ctx);
    JS_SetPropertyStr(ctx, export, "Matrix", matrix_ctor);
    JS_SetPropertyStr(ctx, export, "Array", array_ctor);
    return export;
}