import t, c
from t import CInt, CLong, CDouble, CPtr, CSizeT
from stdio import printf
import cpython


_pass_count: CInt = 0
_fail_count: CInt = 0

def check(name: t.CChar | CPtr, condition: bool):
    global _pass_count, _fail_count
    if condition:
        _pass_count += 1
        printf("  [PASS] %s\n", name)
    else:
        _fail_count += 1
        printf("  [FAIL] %s\n", name)


def test_init_finalize():
    printf("\n+-- Py_Initialize / Py_Finalize --+\n")
    cpython.Py_Initialize()
    check("Py_Initialize succeeded", cpython.Py_IsInitialized() != 0)
    cpython.Py_Finalize()
    check("Py_Finalize succeeded", cpython.Py_IsInitialized() == 0)
    printf("+--------------------------------------------+\n")


def test_long_operations():
    printf("\n+-- PyLong operations --+\n")
    cpython.Py_Initialize()

    # 创建整数
    py_num: cpython.PyObject | CPtr = cpython.PyLong_FromLong(12345)
    check("PyLong_FromLong != NULL", py_num != 0)

    # 整数转字符串
    py_str: cpython.PyObject | CPtr = cpython.PyObject_Str(py_num)
    check("PyObject_Str != NULL", py_str != 0)

    # 获取 C 字符串
    c_str: t.CChar | CPtr = cpython.PyUnicode_AsUTF8(py_str)
    check("PyUnicode_AsUTF8 != NULL", c_str != 0)
    printf("  PyLong_FromLong(12345) -> str = %s\n", c_str)

    # 读取回来
    val: CLong = cpython.PyLong_AsLong(py_num)
    check("PyLong_AsLong == 12345", val == 12345)

    # 负数
    py_neg: cpython.PyObject | CPtr = cpython.PyLong_FromLong(-999)
    neg_val: CLong = cpython.PyLong_AsLong(py_neg)
    check("PyLong_FromLong(-999) == -999", neg_val == -999)

    # 浮点数转整数
    py_dbl: cpython.PyObject | CPtr = cpython.PyFloat_FromDouble(3.14159)
    check("PyFloat_FromDouble != NULL", py_dbl != 0)
    dbl_val: CDouble = cpython.PyFloat_AsDouble(py_dbl)
    check("PyFloat_AsDouble ~= 3.14", dbl_val > 3.14 and dbl_val < 3.15)

    # 清理
    cpython.Py_DecRef(py_dbl)
    cpython.Py_DecRef(py_neg)
    cpython.Py_DecRef(py_str)
    cpython.Py_DecRef(py_num)

    cpython.Py_Finalize()
    printf("+--------------------------------------------+\n")


def test_string_operations():
    printf("\n+-- PyUnicode operations --+\n")
    cpython.Py_Initialize()

    # 从 C 字符串创建
    py_s1: cpython.PyObject | CPtr = cpython.PyUnicode_FromString("Hello, Viper!")
    check("PyUnicode_FromString != NULL", py_s1 != 0)

    # 获取 UTF8
    utf8: t.CChar | CPtr = cpython.PyUnicode_AsUTF8(py_s1)
    check("AsUTF8 != NULL", utf8 != 0)
    printf("  PyUnicode_FromString -> %s\n", utf8)

    # 字符串长度
    length: CLong = cpython.PyUnicode_GetLength(py_s1)
    check("GetLength == 13", length == 13)

    # 字符串拼接
    py_s2: cpython.PyObject | CPtr = cpython.PyUnicode_FromString(" World")
    py_concat: cpython.PyObject | CPtr = cpython.PyUnicode_Concat(py_s1, py_s2)
    concat_utf8: t.CChar | CPtr = cpython.PyUnicode_AsUTF8(py_concat)
    check("Concat != NULL", py_concat != 0)
    printf("  Concat -> %s\n", concat_utf8)

    # 清理
    cpython.Py_DecRef(py_concat)
    cpython.Py_DecRef(py_s2)
    cpython.Py_DecRef(py_s1)

    cpython.Py_Finalize()
    printf("+--------------------------------------------+\n")


def test_list_operations():
    printf("\n+-- PyList operations --+\n")
    cpython.Py_Initialize()

    # 创建列表
    py_list: cpython.PyObject | CPtr = cpython.PyList_New(0)
    check("PyList_New != NULL", py_list != 0)

    # Append 元素
    py_i1: cpython.PyObject | CPtr = cpython.PyLong_FromLong(10)
    py_i2: cpython.PyObject | CPtr = cpython.PyLong_FromLong(20)
    py_i3: cpython.PyObject | CPtr = cpython.PyLong_FromLong(30)
    cpython.PyList_Append(py_list, py_i1)
    cpython.PyList_Append(py_list, py_i2)
    cpython.PyList_Append(py_list, py_i3)

    # 检查大小
    list_size: CSizeT = cpython.PyList_Size(py_list)
    check("PyList_Size == 3", list_size == 3)

    # 获取元素
    item0: cpython.PyObject | CPtr = cpython.PyList_GetItem(py_list, 0)
    val0: CLong = cpython.PyLong_AsLong(item0)
    check("list[0] == 10", val0 == 10)

    item2: cpython.PyObject | CPtr = cpython.PyList_GetItem(py_list, 2)
    val2: CLong = cpython.PyLong_AsLong(item2)
    check("list[2] == 30", val2 == 30)

    # 清理
    cpython.Py_DecRef(py_i3)
    cpython.Py_DecRef(py_i2)
    cpython.Py_DecRef(py_i1)
    cpython.Py_DecRef(py_list)

    cpython.Py_Finalize()
    printf("+--------------------------------------------+\n")


def test_dict_operations():
    printf("\n+-- PyDict operations --+\n")
    cpython.Py_Initialize()

    # 创建字典
    py_dict: cpython.PyObject | CPtr = cpython.PyDict_New()
    check("PyDict_New != NULL", py_dict != 0)

    # 设置键值对
    py_key1: cpython.PyObject | CPtr = cpython.PyUnicode_FromString("name")
    py_val1: cpython.PyObject | CPtr = cpython.PyUnicode_FromString("Viper")
    cpython.PyDict_SetItem(py_dict, py_key1, py_val1)

    py_key2: cpython.PyObject | CPtr = cpython.PyUnicode_FromString("version")
    py_val2: cpython.PyObject | CPtr = cpython.PyLong_FromLong(1)
    cpython.PyDict_SetItem(py_dict, py_key2, py_val2)

    # 检查大小
    dict_size: CSizeT = cpython.PyDict_Size(py_dict)
    check("PyDict_Size == 2", dict_size == 2)

    # 获取值
    got_val: cpython.PyObject | CPtr = cpython.PyDict_GetItem(py_dict, py_key1)
    check("GetItem(name) != NULL", got_val != 0)
    got_str: t.CChar | CPtr = cpython.PyUnicode_AsUTF8(got_val)
    check("name == Viper", got_str != 0)
    printf("  dict['name'] = %s\n", got_str)

    # 使用 GetItemString
    got_ver: cpython.PyObject | CPtr = cpython.PyDict_GetItemString(py_dict, "version")
    check("GetItemString(version) != NULL", got_ver != 0)
    ver_val: CLong = cpython.PyLong_AsLong(got_ver)
    check("version == 1", ver_val == 1)

    # 清理
    cpython.Py_DecRef(py_val2)
    cpython.Py_DecRef(py_key2)
    cpython.Py_DecRef(py_val1)
    cpython.Py_DecRef(py_key1)
    cpython.Py_DecRef(py_dict)

    cpython.Py_Finalize()
    printf("+--------------------------------------------+\n")


def test_pyrun_string():
    printf("\n+-- PyRun_SimpleString --+\n")
    cpython.Py_Initialize()

    # 执行简单 Python 代码
    rc: CInt = cpython.PyRun_SimpleString("print('Hello from embedded Python!')")
    check("PyRun_SimpleString == 0", rc == 0)

    # 执行赋值
    rc2: CInt = cpython.PyRun_SimpleString("x = 6 * 7")
    check("PyRun_SimpleString(assign) == 0", rc2 == 0)

    # 读取结果
    rc3: CInt = cpython.PyRun_SimpleString("print(f'The answer is {x}')")
    check("PyRun_SimpleString(print) == 0", rc3 == 0)

    cpython.Py_Finalize()
    printf("+--------------------------------------------+\n")


def test_import_module():
    printf("\n+-- PyImport_ImportModule --+\n")
    cpython.Py_Initialize()

    # 导入 math 模块
    math_mod: cpython.PyObject | CPtr = cpython.PyImport_ImportModule("math")
    check("import math != NULL", math_mod != 0)

    if math_mod != 0:
        # 获取 math.pi
        pi_val: cpython.PyObject | CPtr = cpython.PyObject_GetAttrString(math_mod, "pi")
        check("math.pi != NULL", pi_val != 0)
        if pi_val != 0:
            pi_dbl: CDouble = cpython.PyFloat_AsDouble(pi_val)
            check("math.pi ~= 3.14", pi_dbl > 3.14 and pi_dbl < 3.15)
            printf("  math.pi = %f\n", pi_dbl)
            cpython.Py_DecRef(pi_val)

        # 获取 math.sqrt 函数
        sqrt_func: cpython.PyObject | CPtr = cpython.PyObject_GetAttrString(math_mod, "sqrt")
        check("math.sqrt != NULL", sqrt_func != 0)
        if sqrt_func != 0:
            cpython.Py_DecRef(sqrt_func)

        cpython.Py_DecRef(math_mod)

    cpython.Py_Finalize()
    printf("+--------------------------------------------+\n")


def test_operator_overload():
    printf("\n+-- Operator Overload (PyLongObj / PyFloatObj / PyUnicodeObj / PyListObj / PyDictObj) --+\n")
    cpython.Py_Initialize()

    # --- PyLongObj: from_long, to_long, +, -, *, - (neg) ---
    a: cpython.PyLongObj | CPtr = cpython.PyLongObj.from_long(10)
    b: cpython.PyLongObj | CPtr = cpython.PyLongObj.from_long(3)
    check("PyLongObj.from_long != NULL", a != 0 and b != 0)

    add_r: cpython.PyLongObj | CPtr = a + b
    check("PyLongObj 10 + 3 == 13", add_r.to_long() == 13)

    sub_r: cpython.PyLongObj | CPtr = a - b
    check("PyLongObj 10 - 3 == 7", sub_r.to_long() == 7)

    mul_r: cpython.PyLongObj | CPtr = a * b
    check("PyLongObj 10 * 3 == 30", mul_r.to_long() == 30)

    neg_r: cpython.PyLongObj | CPtr = -b
    check("PyLongObj -3 == -3", neg_r.to_long() == -3)

    mul_r.dec_ref()
    sub_r.dec_ref()
    add_r.dec_ref()
    neg_r.dec_ref()
    b.dec_ref()
    a.dec_ref()

    # --- PyFloatObj: from_double, to_double, +, -, *, /, - (neg) ---
    fa: cpython.PyFloatObj | CPtr = cpython.PyFloatObj.from_double(10.0)
    fb: cpython.PyFloatObj | CPtr = cpython.PyFloatObj.from_double(4.0)
    check("PyFloatObj.from_double != NULL", fa != 0 and fb != 0)

    fadd: cpython.PyFloatObj | CPtr = fa + fb
    check("PyFloatObj 10.0 + 4.0 == 14.0", fadd.to_double() == 14.0)

    fsub: cpython.PyFloatObj | CPtr = fa - fb
    check("PyFloatObj 10.0 - 4.0 == 6.0", fsub.to_double() == 6.0)

    fmul: cpython.PyFloatObj | CPtr = fa * fb
    check("PyFloatObj 10.0 * 4.0 == 40.0", fmul.to_double() == 40.0)

    fdiv: cpython.PyFloatObj | CPtr = fa / fb
    check("PyFloatObj 10.0 / 4.0 == 2.5", fdiv.to_double() == 2.5)

    fneg: cpython.PyFloatObj | CPtr = -fb
    check("PyFloatObj -4.0 == -4.0", fneg.to_double() == -4.0)

    fadd.dec_ref()
    fsub.dec_ref()
    fmul.dec_ref()
    fdiv.dec_ref()
    fneg.dec_ref()
    fb.dec_ref()
    fa.dec_ref()

    # --- PyUnicodeObj: from_string, as_utf8, len, + (concat) ---
    sa: cpython.PyUnicodeObj | CPtr = cpython.PyUnicodeObj.from_string("Hello")
    sb: cpython.PyUnicodeObj | CPtr = cpython.PyUnicodeObj.from_string(" World")
    check("PyUnicodeObj.from_string != NULL", sa != 0 and sb != 0)

    sc: cpython.PyUnicodeObj | CPtr = sa + sb
    sc_utf8: t.CChar | CPtr = sc.as_utf8()
    check("PyUnicodeObj concat != NULL", sc_utf8 != 0)
    if sc_utf8 != 0:
        printf("  PyUnicodeObj 'Hello' + ' World' = %s\n", sc_utf8)

    sa_len: t.CPtrDiffT = len(sa)
    check("PyUnicodeObj len('Hello') == 5", sa_len == 5)

    sc.dec_ref()
    sb.dec_ref()
    sa.dec_ref()

    # --- PyListObj: new, append, len, [], []= ---
    lst: cpython.PyListObj | CPtr = cpython.PyListObj.new()
    check("PyListObj.new != NULL", lst != 0)
    if lst != 0:
        e1: cpython.PyObject | CPtr = cpython.PyLong_FromLong(100)
        e2: cpython.PyObject | CPtr = cpython.PyLong_FromLong(200)
        e3: cpython.PyObject | CPtr = cpython.PyLong_FromLong(300)
        lst.append(e1)
        lst.append(e2)
        lst.append(e3)
        check("PyListObj len == 3", len(lst) == 3)

        item0: cpython.PyObject | CPtr = lst[t.CSizeT(0)]
        check("PyListObj lst[0] == 100", cpython.PyLong_AsLong(item0) == 100)

        # __setitem__
        e4: cpython.PyObject | CPtr = cpython.PyLong_FromLong(999)
        lst[t.CSizeT(1)] = e4
        item1: cpython.PyObject | CPtr = lst[t.CSizeT(1)]
        check("PyListObj lst[1] = 999", cpython.PyLong_AsLong(item1) == 999)

        cpython.Py_DecRef(e4)
        cpython.Py_DecRef(e3)
        cpython.Py_DecRef(e2)
        cpython.Py_DecRef(e1)
        lst.dec_ref()

    # --- PyDictObj: new, [], []=, len, set_string, get_string ---
    dct: cpython.PyDictObj | CPtr = cpython.PyDictObj.new()
    check("PyDictObj.new != NULL", dct != 0)
    if dct != 0:
        dk1: cpython.PyObject | CPtr = cpython.PyUnicode_FromString("lang")
        dv1: cpython.PyObject | CPtr = cpython.PyUnicode_FromString("Viper")
        dct[dk1] = dv1
        check("PyDictObj len == 1", len(dct) == 1)

        got_lang: cpython.PyObject | CPtr = dct[dk1]
        got_str: t.CChar | CPtr = cpython.PyUnicode_AsUTF8(got_lang)
        check("PyDictObj dct['lang'] == Viper", got_str != 0)
        if got_str != 0:
            printf("  PyDictObj dct['lang'] = %s\n", got_str)

        # get_string / set_string
        dct.set_string("ver", cpython.PyLong_FromLong(3))
        got_ver: cpython.PyObject | CPtr = dct.get_string("ver")
        check("PyDictObj get_string('ver') == 3", got_ver != 0 and cpython.PyLong_AsLong(got_ver) == 3)

        cpython.Py_DecRef(dv1)
        cpython.Py_DecRef(dk1)
        dct.dec_ref()

    cpython.Py_Finalize()
    printf("+--------------------------------------------+\n")


def main() -> CInt | t.CExport:
    printf("=== CPython Integration Test ===\n")

    test_init_finalize()
    test_long_operations()
    test_string_operations()
    test_list_operations()
    test_dict_operations()
    test_pyrun_string()
    test_import_module()
    test_operator_overload()

    printf("\n=== Results: %d passed, %d failed ===\n", _pass_count, _fail_count)
    return 0