from stdlib import malloc, free
from stdint import *
import string
import w32.win32console
from stdarg import va_start, va_end, va_arg, va_list
import t
import c


class Struct1:
    a: UINT
    b: UINT = 123


class Item:
    data: t.CVoid | t.CPtr
    size: t.CSizeT


class List:
    items: Item | t.CPtr
    count: t.CSizeT = 0
    capacity: t.CSizeT = 4

    def __init__(self):
        self.count = 0
        self.capacity = 4
        self.items = malloc(self.capacity * Item.__sizeof__())

    def append(self, data: t.CVoid | t.CPtr, size: t.CSizeT):
        self.count += 1
        item: Item | t.CPtr = self.items[self.count - 1]
        item.size = size
        item.data = malloc(size + 1)
        memcpy(item.data, data, size)

    def get(self, index: t.CSizeT) -> t.CInt8T | t.CPtr:
        if index >= self.count:
            return None
        return self.items[index].data

    def free(self):
        i: t.CSizeT = 0
        for i in range(self.count):
            if self.items[i].data != 0:
                free(self.items[i].data)
        free(self.items)


class EnumTest(t.CEnum):
    A: c.State
    B: c.State
    C: c.State
    Len: c.State


@t.Object
class OOPTest:
    a: UINT
    b: UINT = 12345
    c: UINT

    def __init__(self):
        self.a = 451

    def set_c(self, v: UINT):
        self.c = v

    def __call__(self, a: UINT, b: UINT) -> UINT:
        return a + b

    def __enter__(self) -> 'OOPTest' | t.CPtr:
        return self

    def __exit__(self):
        self.a = 99
        self.b = 99
        self.c = 99

    def __add__(self, other: 'OOPTest' | t.CPtr) -> 'OOPTest' | t.CPtr:
        s = OOPTest()
        s.a = (self.a + other.a)
        return s

    def __str__(self) -> str:
        return f"OOPTest ({self.a}) ({self.b}) ({self.c})"

    def __bool__(self) -> bool:
        return bool(self.a != 0)

    def __len__(self) -> UINT:
        return self.a


@t.Object
class RangeOOPSimulation:
    start: UINT
    end: UINT

    def __init__(self, start: UINT, end: UINT):
        self.start = start
        self.end = end

    def __iter__(self) -> 'RangeOOPSimulation' | t.CPtr:
        return self

    def __next__(self) -> UINT:
        if self.start >= self.end:
            raise StopIteration
        self.start += 1
        return self.start - 1


GLOBAL_VAR: t.CInt = 0


def rebutnotreGLOBAL_VAR() -> t.CInt:
    GLOBAL_VAR = 100


def reGLOBAL_VAR() -> t.CInt:
    global GLOBAL_VAR
    GLOBAL_VAR = 100


def function2function() -> t.CInt:
    def func() -> t.CInt:
        x = 1000

        def func2() -> t.CInt:
            nonlocal x
            x = 1000000

        func2()
        return x

    return func()


def manyreturn() -> tuple[t.CInt, t.CInt]:
    return 114, 514


def str_to_hex(s: str) -> str:
    buf: str = malloc(1024)
    out: str = buf
    hex_chars: list[t.CChar, None] = "0123456789ABCDEF"
    while c.Cast(s) != 0:
        cr: t.CUnsignedChar = c.CCast(s)
        c.Set(c.CCast(out), hex_chars[(cr >> 4) & 0xF])
        out += 1
        c.Set(c.CCast(out), hex_chars[cr & 0xF])
        out += 1
        s += 1
    c.Set(c.CCast(out), 0)
    return buf


def int_to_hex(num: int) -> str:
    buf = t.CChar(malloc(128), t.CPtr)
    if not buf:
        return None

    hex_chars: list[t.CChar, None] = "0123456789ABCDEF"
    i = 0

    buf[i] = '0'
    buf[i + 1] = 'x'
    i += 2

    if num == 0:
        buf[i] = '0'
        i += 1
        buf[i] = '0'
        i += 1
        buf[i] = '\0'
        return buf

    temp: list[t.CChar, 32]
    temp_idx = 0

    n: t.CUnsignedInt = t.CUnsignedInt(num)

    while n > 0:
        rem: t.CInt = n & 0xF
        temp[temp_idx] = hex_chars[rem]
        temp_idx += 1
        n = n >> 4

    while temp_idx > 0:
        temp_idx -= 1
        buf[i] = temp[temp_idx]
        i += 1

    buf[i] = '\0'
    return buf


def int_to_str(num: int) -> str:
    buf: t.CChar | t.CPtr = t.CChar(malloc(128), t.CPtr)
    if not buf:
        return None
    is_negative = 0
    i = 0
    if num == 0:
        buf[i] = '0'
        i += 1
        buf[i] = '\0'
        return buf
    if num < 0:
        is_negative = 1
        num = -num
    while num > 0:
        rem = num % 10
        buf[i] = '0' + rem
        i += 1
        num = num / 10
    if is_negative:
        buf[i] = '-'
        i += 1
    buf[i] = '\0'
    start = 0
    end = i - 1
    while start < end:
        tr: t.CChar = buf[start]
        buf[start] = buf[end]
        buf[end] = tr
        start += 1
        end -= 1
    return buf


def main():
    w32.win32console.SetConsoleOutputCP(65001)
    w32.win32console.SetConsoleCP(65001)

    print("============== BENCHMARK 1: Structures & OOP ==============")

    print("\nTesting Struct1:")
    s1_: Struct1 = Struct1(10)
    print("s1_.a:", s1_.a)
    print("s1_.b:", s1_.b)

    print("\nTesting OOPTest class:")
    d = OOPTest()
    d.set_c(100)
    d.a = 1
    print("d.a:", d.a)
    print("d.b:", d.b)
    print("d.c:", d.c)

    print("\nTesting __call__ (d(1, 2)):")
    print("d(1, 2):", d(1, 2))

    print("\nTesting __bool__:")
    if d:
        print("d is truthy (d.a != 0)")
    else:
        print("d is falsy (d.a == 0)")

    print("\nTesting __len__:")
    print("len(d):", len(d))

    print("\nTesting __str__:")
    print("str(d):", str(d))

    print("\nTesting __add__:")
    l = OOPTest()
    l.a = 100
    result = d + l
    print("(d + l).a:", result.a)
    free(l)
    free(result)

    print("\nTesting context manager (with statement):")
    with OOPTest() as d2:
        print("Inside with block, d2(1, 2):", d2(1, 2))
    print("After exit, d2.__str__():", str(d2))

    print("\nTesting List class:")
    l = List()
    nihao = "你好！世界"
    l.append(c.Addr(114514), t.CInt.__sizeof__())
    l.append(nihao, len(nihao))
    print("List count:", l.count)
    l.free()

    print("\n============== BENCHMARK 2: Global Variables & Functions ==============")

    print("\nTesting global variable re-assignment:")
    print("Initial GLOBAL_VAR:", GLOBAL_VAR)
    rebutnotreGLOBAL_VAR()
    print("After rebutnotreGLOBAL_VAR():", GLOBAL_VAR)
    reGLOBAL_VAR()
    print("After reGLOBAL_VAR():", GLOBAL_VAR)

    print("\nTesting function2function with nonlocal:")
    print("function2function():", function2function())

    print("\nTesting tuple return:")
    a: t.CInt
    b: t.CInt
    a, b = manyreturn()
    print("manyreturn():", a, b)

    print("\nTesting walrus operator:")
    kkkkk = 0
    if not GLOBAL_VAR or (kkkkk := 1):
        print("GLOBAL_VAR is 0, kkkkk:", kkkkk)
    print("Final kkkkk:", kkkkk)

    print("\n============== BENCHMARK 3: Memory & Pointers ==============")

    print("\nTesting malloc/free:")
    p2 = malloc(100)
    i: t.CInt = 0
    while i <= 100:
        p2[i] = 4
        i += 1
    all_four = 1
    for i in range(0, 100 + 1):
        if p2[i] != 4:
            all_four = 0
            break
    if all_four:
        print("All 4 - Memory write/read verified")
    free(p2)

    print("\nTesting for loop with range:")
    add_num = 0
    for i in range(0, 100 + 1):
        add_num += 1
    print("i:", i, "add_num:", add_num)

    print("\nTesting for loop with continue:")
    lll = 0
    for i in range(0, 10 + 1):
        if i % 2 != 0:
            continue
        lll += i
    print("Sum of even numbers 0-10:", lll)

    print("\n============== BENCHMARK 4: Iteration & Strings ==============")

    print("\nTesting RangeOOPSimulation (custom iterator):")
    print("Iterating RangeOOPSimulation(0, 6):")
    for i in RangeOOPSimulation(0, 6):
        print("RangeOOPSimulation item:", i)

    print("\nTesting str_to_hex:")
    hex_str = str_to_hex("ABCDEF")
    print("hex of ABCDEF:", hex_str)
    free(hex_str)

    print("\nTesting int_to_hex:")
    hex_num = int_to_hex(255)
    print("hex of 255:", hex_num)
    free(hex_num)

    print("\nTesting int_to_str:")
    num_str = int_to_str(-12345)
    print("str of -12345:", num_str)
    free(num_str)

    print("\nTesting f-string with Chinese characters:")
    chinese = "你好"
    print(f"{chinese}，世界")

    print("\n============== BENCHMARK 5: Exception Handling ==============")

    print("\nTesting try/except/else/finally:")
    try:
        print("In try block")
        assert 1 != 1, "1 == 1"
        print("This should not print")
    except AssertionError as e:
        print("AssertionError caught:", e)
    except ValueError as e:
        print("ValueError caught:", e)
    else:
        print("No exception in try block")
    finally:
        print("Finally block executed")

    free(d)
    free(d2)

    print("\n============== All Benchmarks Completed ==============")


main()