#!/usr/bin/env python3
import os
from test_limbo_cli import TestLimboShell

sqlite_exec = "./target/debug/limbo"
sqlite_flags = os.getenv("SQLITE_FLAGS", "-q").split(" ")

test_data = """CREATE TABLE numbers ( id INTEGER PRIMARY KEY, value FLOAT NOT NULL);
INSERT INTO numbers (value) VALUES (1.0);
INSERT INTO numbers (value) VALUES (2.0);
INSERT INTO numbers (value) VALUES (3.0);
INSERT INTO numbers (value) VALUES (4.0);
INSERT INTO numbers (value) VALUES (5.0);
INSERT INTO numbers (value) VALUES (6.0);
INSERT INTO numbers (value) VALUES (7.0);
CREATE TABLE test (value REAL, percent REAL);
INSERT INTO test values (10, 25);
INSERT INTO test values (20, 25);
INSERT INTO test values (30, 25);
INSERT INTO test values (40, 25);
INSERT INTO test values (50, 25);
INSERT INTO test values (60, 25);
INSERT INTO test values (70, 25);
"""


def validate_string_uuid(res):
    return len(res) == 36


def test_uuid():
    limbo = TestLimboShell()
    specific_time = "01945ca0-3189-76c0-9a8f-caf310fc8b8e"
    # these are built into the binary, so we just test they work
    limbo.run_test_fn(
        "SELECT hex(uuid4());",
        lambda res: int(res, 16) is not None,
        "uuid functions are registered properly with ext loaded",
    )
    limbo.run_test_fn("SELECT uuid4_str();", lambda res: len(res) == 36)
    limbo.run_test_fn("SELECT hex(uuid7());", lambda res: int(res, 16) is not None)
    limbo.run_test_fn(
        "SELECT uuid7_timestamp_ms(uuid7()) / 1000;", lambda res: res.isdigit()
    )
    limbo.run_test_fn("SELECT uuid7_str();", validate_string_uuid)
    limbo.run_test_fn("SELECT uuid_str(uuid7());", validate_string_uuid)
    limbo.run_test_fn(
        "SELECT hex(uuid_blob(uuid7_str()));", lambda res: int(res, 16) is not None
    )
    limbo.run_test_fn("SELECT uuid_str(uuid_blob(uuid7_str()));", validate_string_uuid)
    limbo.run_test_fn(
        f"SELECT uuid7_timestamp_ms('{specific_time}') / 1000;",
        lambda res: res == "1736720789",
    )
    limbo.run_test_fn(
        "SELECT gen_random_uuid();",
        validate_string_uuid,
        "scalar alias's are registered properly",
    )
    limbo.quit()


def true(res):
    return res == "1"


def false(res):
    return res == "0"


def null(res):
    return res == ""


def test_regexp():
    limbo = TestLimboShell(test_data)
    extension_path = "./target/debug/liblimbo_regexp"
    # before extension loads, assert no function
    limbo.run_test_fn(
        "SELECT regexp('a.c', 'abc');",
        lambda res: "Parse error: no such function" in res,
    )
    limbo.run_test_fn(f".load {extension_path}", null)
    print(f"Extension {extension_path} loaded successfully.")
    limbo.run_test_fn("SELECT regexp('a.c', 'abc');", true)
    limbo.run_test_fn("SELECT regexp('a.c', 'ac');", false)
    limbo.run_test_fn("SELECT regexp('[0-9]+', 'the year is 2021');", true)
    limbo.run_test_fn("SELECT regexp('[0-9]+', 'the year is unknow');", false)
    limbo.run_test_fn("SELECT regexp_like('the year is 2021', '[0-9]+');", true)
    limbo.run_test_fn("SELECT regexp_like('the year is unknow', '[0-9]+');", false)
    limbo.run_test_fn(
        "SELECT regexp_substr('the year is 2021', '[0-9]+') = '2021';",
        true,
    )
    limbo.run_test_fn("SELECT regexp_substr('the year is unknow', '[0-9]+');", null)
    limbo.run_test_fn(
        "select regexp_replace('the year is 2021', '[0-9]+', '2050') = 'the year is 2050';",
        true,
    )
    limbo.run_test_fn(
        "select regexp_replace('the year is 2021', '2k21', '2050') = 'the year is 2021';",
        true,
    )
    limbo.run_test_fn(
        "select regexp_replace('the year is 2021', '([0-9]+)', '$1 or 2050') = 'the year is 2021 or 2050';",
        true,
    )
    limbo.quit()


def validate_median(res):
    return res == "4.0"


def validate_median_odd(res):
    return res == "4.5"


def validate_percentile1(res):
    return res == "25.0"


def validate_percentile2(res):
    return res == "43.0"


def validate_percentile_disc(res):
    return res == "40.0"


def test_aggregates():
    limbo = TestLimboShell(init_commands=test_data)
    extension_path = "./target/debug/liblimbo_percentile"
    # assert no function before extension loads
    limbo.run_test_fn(
        "SELECT median(1);",
        lambda res: "error: no such function: " in res,
        "median agg function returns null when ext not loaded",
    )
    limbo.execute_dot(f".load {extension_path}")
    limbo.run_test_fn(
        "select median(value) from numbers;",
        validate_median,
        "median agg function works",
    )
    limbo.execute_dot("INSERT INTO numbers (value) VALUES (8.0);\n")
    limbo.run_test_fn(
        "select median(value) from numbers;",
        validate_median_odd,
        "median agg function works with odd number of elements",
    )
    limbo.run_test_fn(
        "SELECT percentile(value, percent) from test;",
        validate_percentile1,
        "test aggregate percentile function with 2 arguments works",
    )
    limbo.run_test_fn(
        "SELECT percentile(value, 55) from test;",
        validate_percentile2,
        "test aggregate percentile function with 1 argument works",
    )
    limbo.run_test_fn(
        "SELECT percentile_cont(value, 0.25) from test;", validate_percentile1
    )
    limbo.run_test_fn(
        "SELECT percentile_disc(value, 0.55) from test;", validate_percentile_disc
    )
    limbo.quit()


# Encoders and decoders
def validate_url_encode(a):
    return a == "%2Fhello%3Ftext%3D%28%E0%B2%A0_%E0%B2%A0%29"


def validate_url_decode(a):
    return a == "/hello?text=(ಠ_ಠ)"


def validate_hex_encode(a):
    return a == "68656c6c6f"


def validate_hex_decode(a):
    return a == "hello"


def validate_base85_encode(a):
    return a == "BOu!rDZ"


def validate_base85_decode(a):
    return a == "hello"


def validate_base32_encode(a):
    return a == "NBSWY3DP"


def validate_base32_decode(a):
    return a == "hello"


def validate_base64_encode(a):
    return a == "aGVsbG8="


def validate_base64_decode(a):
    return a == "hello"


def test_crypto():
    limbo = TestLimboShell()
    extension_path = "./target/debug/liblimbo_crypto"
    # assert no function before extension loads
    limbo.run_test_fn(
        "SELECT crypto_blake('a');",
        lambda res: "Parse error" in res,
        "crypto_blake3 returns null when ext not loaded",
    )
    limbo.execute_dot(f".load {extension_path}")
    # Hashing and Decode
    limbo.run_test_fn(
        "SELECT crypto_encode(crypto_blake3('abc'), 'hex');",
        lambda res: res
        == "6437b3ac38465133ffb63b75273a8db548c558465d79db03fd359c6cd5bd9d85",
        "blake3 should encrypt correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode(crypto_md5('abc'), 'hex');",
        lambda res: res == "900150983cd24fb0d6963f7d28e17f72",
        "md5 should encrypt correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode(crypto_sha1('abc'), 'hex');",
        lambda res: res == "a9993e364706816aba3e25717850c26c9cd0d89d",
        "sha1 should encrypt correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode(crypto_sha256('abc'), 'hex');",
        lambda a: a
        == "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
        "sha256 should encrypt correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode(crypto_sha384('abc'), 'hex');",
        lambda a: a
        == "cb00753f45a35e8bb5a03d699ac65007272c32ab0eded1631a8b605a43ff5bed8086072ba1e7cc2358baeca134c825a7",
        "sha384 should encrypt correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode(crypto_sha512('abc'), 'hex');",
        lambda a: a
        == "ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f",
        "sha512 should encrypt correctly",
    )

    # Encoding and Decoding
    limbo.run_test_fn(
        "SELECT crypto_encode('hello', 'base32');",
        validate_base32_encode,
        "base32 should encode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_decode('NBSWY3DP', 'base32');",
        validate_base32_decode,
        "base32 should decode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode('hello', 'base64');",
        validate_base64_encode,
        "base64 should encode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_decode('aGVsbG8=', 'base64');",
        validate_base64_decode,
        "base64 should decode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode('hello', 'base85');",
        validate_base85_encode,
        "base85 should encode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_decode('BOu!rDZ', 'base85');",
        validate_base85_decode,
        "base85 should decode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode('hello', 'hex');",
        validate_hex_encode,
        "hex should encode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_decode('68656c6c6f', 'hex');",
        validate_hex_decode,
        "hex should decode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_encode('/hello?text=(ಠ_ಠ)', 'url');",
        validate_url_encode,
        "url should encode correctly",
    )
    limbo.run_test_fn(
        "SELECT crypto_decode('%2Fhello%3Ftext%3D%28%E0%B2%A0_%E0%B2%A0%29', 'url');",
        validate_url_decode,
        "url should decode correctly",
    )
    limbo.quit()


def test_series():
    limbo = TestLimboShell()
    ext_path = "./target/debug/liblimbo_series"
    limbo.run_test_fn(
        "SELECT * FROM generate_series(1, 10);",
        lambda res: "Virtual table module not found: generate_series" in res,
    )
    limbo.execute_dot(f".load {ext_path}")
    limbo.run_test_fn(
        "SELECT * FROM generate_series(1, 10);",
        lambda res: res == "1\n2\n3\n4\n5\n6\n7\n8\n9\n10",
    )
    limbo.run_test_fn(
        "SELECT * FROM generate_series(1, 10, 2);",
        lambda res: res == "1\n3\n5\n7\n9",
    )
    limbo.run_test_fn(
        "SELECT * FROM generate_series(1, 10, 2, 3);",
        lambda res: "Invalid Argument" in res,
    )
    limbo.run_test_fn(
        "SELECT * FROM generate_series(10, 1, -2);",
        lambda res: res == "10\n8\n6\n4\n2",
    )
    limbo.quit()


def test_kv():
    ext_path = "./target/debug/liblimbo_kv"
    limbo = TestLimboShell()
    limbo.run_test_fn(
        "create virtual table t using kv_store;",
        lambda res: "Virtual table module not found: kv_store" in res,
    )
    limbo.execute_dot(f".load {ext_path}")
    limbo.run_test_fn(
        "create virtual table t using kv_store;",
        null,
        "can create kv_store vtable",
    )
    limbo.run_test_fn(
        "insert into t values ('hello', 'world');",
        null,
        "can insert into kv_store vtable",
    )
    limbo.run_test_fn(
        "select value from t where key = 'hello';",
        lambda res: "world" == res,
        "can select from kv_store",
    )
    limbo.run_test_fn(
        "delete from t where key = 'hello';",
        null,
        "can delete from kv_store",
    )
    limbo.run_test_fn("insert into t values ('other', 'value');", null)
    limbo.run_test_fn(
        "select value from t where key = 'hello';",
        lambda res: "" == res,
        "proper data is deleted",
    )
    limbo.run_test_fn(
        "select * from t;",
        lambda res: "other|value" == res,
        "can select after deletion",
    )
    limbo.run_test_fn(
        "delete from t where key = 'other';",
        null,
        "can delete from kv_store",
    )
    limbo.run_test_fn(
        "select * from t;",
        lambda res: "" == res,
        "can select empty table without error",
    )
    limbo.run_test_fn(
        "delete from t;",
        null,
        "can delete from empty table without error",
    )
    for i in range(100):
        limbo.execute_dot(f"insert into t values ('key{i}', 'val{i}');")
    limbo.run_test_fn(
        "select count(*) from t;", lambda res: "100" == res, "can insert 100 rows"
    )
    limbo.run_test_fn("delete from t limit 96;", null, "can delete 96 rows")
    limbo.run_test_fn(
        "select count(*) from t;", lambda res: "4" == res, "four rows remain"
    )
    limbo.quit()


if __name__ == "__main__":
    try:
        test_regexp()
        test_uuid()
        test_aggregates()
        test_crypto()
        test_series()
        test_kv()
    except Exception as e:
        print(f"Test FAILED: {e}")
        exit(1)
    print("All tests passed successfully.")