#include "cell.h"
#include <stdint.h>
#include <limits.h>

#define FIT_BITS 56
#define FIT_MAX ((1LL << 55) - 1)
#define FIT_MIN (-(1LL << 55))
#define FIT_MASK ((1ULL << 56) - 1)

static int is_fit(int64_t n) {
    return n >= FIT_MIN && n <= FIT_MAX;
}

static int64_t to_fit_int(JSContext *js, JSValue val) {
    if (!JS_IsNumber(val)) return LLONG_MAX;
    
    int64_t n;
    if (JS_ToInt64(js, &n, val) < 0) return LLONG_MAX;
    
    if (!is_fit(n)) return LLONG_MAX;
    
    return n;
}

static JSValue fit_result(JSContext *js, int64_t result) {
    if (!is_fit(result)) return JS_NULL;
    return JS_NewInt64(js, result);
}

JSC_CCALL(fit_and,
    int64_t first = to_fit_int(js, argv[0]);
    int64_t second = to_fit_int(js, argv[1]);
    
    if (first == LLONG_MAX || second == LLONG_MAX) return JS_NULL;
    
    int64_t result = first & second;
    return fit_result(js, result);
)

JSC_CCALL(fit_left,
    int64_t first = to_fit_int(js, argv[0]);
    int64_t second = to_fit_int(js, argv[1]);
    
    if (first == LLONG_MAX || second == LLONG_MAX) return JS_NULL;
    if (second < 0 || second >= FIT_BITS) return JS_NULL;
    
    // Mask to 56 bits then sign extend
    uint64_t unsigned_first = (uint64_t)first & FIT_MASK;
    uint64_t result = (unsigned_first << second) & FIT_MASK;
    
    // Sign extend from bit 55
    if (result & (1ULL << 55)) {
        result |= ~FIT_MASK;
    }
    
    return JS_NewInt64(js, (int64_t)result);
)

JSC_CCALL(fit_mask,
    if (!JS_IsNumber(argv[0])) return JS_NULL;
    
    int32_t n;
    if (JS_ToInt32(js, &n, argv[0]) < 0) return JS_NULL;
    
    if (n > FIT_BITS || n < -FIT_BITS) return JS_NULL;
    
    if (n == 0) return JS_NewInt64(js, 0);
    
    int64_t result;
    if (n > 0) {
        // Create n ones
        if (n == FIT_BITS) {
            result = -1;
        } else {
            result = (1LL << n) - 1;
        }
    } else {
        // Create -n zeros (which is 56+n ones)
        int ones = FIT_BITS + n;
        if (ones == 0) {
            result = 0;
        } else if (ones == FIT_BITS) {
            result = -1;
        } else {
            uint64_t mask = (1ULL << ones) - 1;
            result = ~mask;
            // Ensure result is within 56-bit range
            result = (int64_t)((uint64_t)result & FIT_MASK);
            // Sign extend
            if (result & (1LL << 55)) {
                result |= ~FIT_MASK;
            }
        }
    }
    
    return JS_NewInt64(js, result);
)

JSC_CCALL(fit_not,
    int64_t val = to_fit_int(js, argv[0]);
    if (val == LLONG_MAX) return JS_NULL;
    
    uint64_t result = ~(uint64_t)val & FIT_MASK;
    
    // Sign extend
    if (result & (1ULL << 55)) {
        result |= ~FIT_MASK;
    }
    
    return JS_NewInt64(js, (int64_t)result);
)

JSC_CCALL(fit_ones,
    int64_t val = to_fit_int(js, argv[0]);
    if (val == LLONG_MAX) return JS_NULL;
    
    uint64_t uval = (uint64_t)val & FIT_MASK;
    int count = 0;
    
    for (int i = 0; i < FIT_BITS; i++) {
        if (uval & (1ULL << i)) count++;
    }
    
    return JS_NewInt32(js, count);
)

JSC_CCALL(fit_or,
    int64_t first = to_fit_int(js, argv[0]);
    int64_t second = to_fit_int(js, argv[1]);
    
    if (first == LLONG_MAX || second == LLONG_MAX) return JS_NULL;
    
    int64_t result = first | second;
    return fit_result(js, result);
)

JSC_CCALL(fit_reverse,
    int64_t val = to_fit_int(js, argv[0]);
    if (val == LLONG_MAX) return JS_NULL;
    
    uint64_t uval = (uint64_t)val & FIT_MASK;
    uint64_t result = 0;
    
    for (int i = 0; i < FIT_BITS; i++) {
        if (uval & (1ULL << i)) {
            result |= (1ULL << (FIT_BITS - 1 - i));
        }
    }
    
    // Sign extend
    if (result & (1ULL << 55)) {
        result |= ~FIT_MASK;
    }
    
    return JS_NewInt64(js, (int64_t)result);
)

JSC_CCALL(fit_right,
    int64_t first = to_fit_int(js, argv[0]);
    int64_t second = to_fit_int(js, argv[1]);
    
    if (first == LLONG_MAX || second == LLONG_MAX) return JS_NULL;
    if (second < 0 || second >= FIT_BITS) return JS_NULL;
    
    // Zero fill right shift
    uint64_t ufirst = (uint64_t)first & FIT_MASK;
    uint64_t result = ufirst >> second;
    
    return JS_NewInt64(js, (int64_t)result);
)

JSC_CCALL(fit_right_signed,
    int64_t first = to_fit_int(js, argv[0]);
    int64_t second = to_fit_int(js, argv[1]);
    
    if (first == LLONG_MAX || second == LLONG_MAX) return JS_NULL;
    if (second < 0 || second >= FIT_BITS) return JS_NULL;
    
    // Sign extend to 64 bits for arithmetic shift
    int64_t extended = first;
    if (first & (1LL << 55)) {
        extended |= ~FIT_MASK;
    }
    
    int64_t result = extended >> second;
    
    // Ensure result is within fit range
    return fit_result(js, result);
)

JSC_CCALL(fit_rotate,
    int64_t first = to_fit_int(js, argv[0]);
    int64_t second = to_fit_int(js, argv[1]);
    
    if (first == LLONG_MAX || second == LLONG_MAX) return JS_NULL;
    
    // Normalize rotation amount
    int rotation = ((int)second % FIT_BITS + FIT_BITS) % FIT_BITS;
    
    uint64_t ufirst = (uint64_t)first & FIT_MASK;
    uint64_t result = ((ufirst << rotation) | (ufirst >> (FIT_BITS - rotation))) & FIT_MASK;
    
    // Sign extend
    if (result & (1ULL << 55)) {
        result |= ~FIT_MASK;
    }
    
    return JS_NewInt64(js, (int64_t)result);
)

JSC_CCALL(fit_xor,
    int64_t first = to_fit_int(js, argv[0]);
    int64_t second = to_fit_int(js, argv[1]);
    
    if (first == LLONG_MAX || second == LLONG_MAX) return JS_NULL;
    
    int64_t result = first ^ second;
    return fit_result(js, result);
)

JSC_CCALL(fit_zeros,
    int64_t val = to_fit_int(js, argv[0]);
    if (val == LLONG_MAX) return JS_NULL;
    
    uint64_t uval = (uint64_t)val & FIT_MASK;
    
    int count = 0;
    for (int i = FIT_BITS - 1; i >= 0; i--) {
        if (uval & (1ULL << i)) break;
        count++;
    }
    
    return JS_NewInt32(js, count);
)

static const JSCFunctionListEntry js_fit_funcs[] = {
    MIST_FUNC_DEF(fit, and, 2),
    MIST_FUNC_DEF(fit, left, 2),
    MIST_FUNC_DEF(fit, mask, 1),
    MIST_FUNC_DEF(fit, not, 1),
    MIST_FUNC_DEF(fit, ones, 1),
    MIST_FUNC_DEF(fit, or, 2),
    MIST_FUNC_DEF(fit, reverse, 1),
    MIST_FUNC_DEF(fit, right, 2),
    MIST_FUNC_DEF(fit, right_signed, 2),
    MIST_FUNC_DEF(fit, rotate, 2),
    MIST_FUNC_DEF(fit, xor, 2),
    MIST_FUNC_DEF(fit, zeros, 1),
};

JSValue js_fit_use(JSContext *js)
{
    JSValue mod = JS_NewObject(js);
    JS_SetPropertyFunctionList(js, mod, js_fit_funcs, countof(js_fit_funcs));
    return mod;
}