import stdio
import stdlib
import memhub
import testcheck


def test_for_in_range():
    testcheck.section("Test 1: for i in range(n)")
    s: int = 0
    for i in range(10):
        s += i
    testcheck.check(s == 45, "range sum OK (45)", "range sum FAILED expect 45")


def test_for_in_str():
    testcheck.section("Test 2: for ch in str")
    cnt: int = 0
    for ch in "hello":
        cnt += 1
    testcheck.check(cnt == 5, "str iter OK (5)", "str iter FAILED expect 5")


def test_for_in_list():
    testcheck.section("Test 3: for v in list[str]")
    arena: bytes = stdlib.malloc(65536)
    bd: memhub.MemBuddy* = memhub.MemBuddy(arena, 65536)
    lst: list[str] = list[str](bd)
    lst.append("x")
    lst.append("y")
    lst.append("z")
    lst.append("w")
    total: int = 0
    for v in lst:
        total += 1
    testcheck.check(total == 4, "list iter OK (4)", "list iter FAILED expect 4")
    stdlib.free(arena)


def test_zip_list_list():
    testcheck.section("Test 4: zip(list, list)")
    arena: bytes = stdlib.malloc(65536)
    bd: memhub.MemBuddy* = memhub.MemBuddy(arena, 65536)
    la: list[str] = list[str](bd)
    la.append("a")
    la.append("b")
    la.append("c")
    lb: list[str] = list[str](bd)
    lb.append("1")
    lb.append("2")
    lb.append("3")
    cnt: int = 0
    for ca, cb in zip(la, lb):
        stdio.printf("zip[%d]: '%s'='%s'\n", cnt, ca, cb)
        cnt += 1
    testcheck.check(cnt == 3, "zip list+list OK (3)", "zip list+list FAILED")
    stdlib.free(arena)


def test_zip_range_list():
    testcheck.section("Test 5: zip(range, list)")
    arena: bytes = stdlib.malloc(65536)
    bd: memhub.MemBuddy* = memhub.MemBuddy(arena, 65536)
    lst: list[str] = list[str](bd)
    lst.append("p")
    lst.append("q")
    lst.append("r")
    cnt: int = 0
    for n, v in zip(range(3), lst):
        stdio.printf("zip range+list[%d]: %d='%s'\n", cnt, n, v)
        cnt += 1
    testcheck.check(cnt == 3, "zip range+list OK (3)", "zip range+list FAILED")
    stdlib.free(arena)


def test_zip_str_str():
    testcheck.section("Test 6: zip(str, str)")
    sa: str = "abc"
    sb: str = "123"
    cnt: int = 0
    for ca, cb in zip(sa, sb):
        stdio.printf("zip str+str[%d]: '%c'%c\n", cnt, ca, cb)
        cnt += 1
    testcheck.check(cnt == 3, "zip str+str OK (3)", "zip str+str FAILED")


def test_zip_range_len_list():
    testcheck.section("Test 7: zip(range(len), list)")
    arena: bytes = stdlib.malloc(65536)
    bd: memhub.MemBuddy* = memhub.MemBuddy(arena, 65536)
    lst: list[str] = list[str](bd)
    lst.append("www")
    lst.append("gvsds")
    lst.append("com")
    cnt: int = 0
    for idx, v in zip(range(len(lst)), lst):
        stdio.printf("lst[%d]='%s'\n", idx, v)
        cnt += 1
    testcheck.check(cnt == 3, "zip(range(len),list) OK (3)", "zip(range(len),list) FAILED")
    stdlib.free(arena)


def test_enumerate_list():
    testcheck.section("Test 8: enumerate(list)")
    arena: bytes = stdlib.malloc(65536)
    bd: memhub.MemBuddy* = memhub.MemBuddy(arena, 65536)
    lst: list[str] = list[str](bd)
    lst.append("alpha")
    lst.append("beta")
    lst.append("gamma")
    cnt: int = 0
    last_idx: int = -1
    for idx, v in enumerate(lst):
        stdio.printf("enum[%d]='%s'\n", idx, v)
        last_idx = idx
        cnt += 1
    testcheck.check(cnt == 3, "enumerate list OK (3)", "enumerate list FAILED")
    testcheck.check(last_idx == 2, "enumerate last_idx OK (2)", "enumerate last_idx FAILED expect 2")
    stdlib.free(arena)


def test_enumerate_str():
    testcheck.section("Test 9: enumerate(str)")
    cnt: int = 0
    last_idx: int = -1
    for idx, ch in enumerate("ABCD"):
        stdio.printf("enum str[%d]='%c'\n", idx, ch)
        last_idx = idx
        cnt += 1
    testcheck.check(cnt == 4, "enumerate str OK (4)", "enumerate str FAILED expect 4")
    testcheck.check(last_idx == 3, "enumerate str last_idx OK (3)", "enumerate str last_idx FAILED expect 3")


def test_enumerate_range():
    testcheck.section("Test 10: enumerate(range)")
    cnt: int = 0
    last_idx: int = -1
    last_val: int = -1
    for idx, v in enumerate(range(5)):
        stdio.printf("enum range[%d]=%d\n", idx, v)
        last_idx = idx
        last_val = v
        cnt += 1
    testcheck.check(cnt == 5, "enumerate range OK (5)", "enumerate range FAILED expect 5")
    testcheck.check(last_idx == 4, "enumerate range last_idx OK (4)", "enumerate range last_idx FAILED expect 4")
    testcheck.check(last_val == 4, "enumerate range last_val OK (4)", "enumerate range last_val FAILED expect 4")


def test_for_else_no_break():
    testcheck.section("Test 11: for-else (迭代耗尽 → 执行 else)")
    total: int = 0
    for i in range(3):
        total += i
    else:
        total += 100
    testcheck.check(total == 103, "for-else no-break OK (103)", "for-else no-break FAILED expect 103")


def test_for_else_break():
    testcheck.section("Test 12: for-else (break → 跳过 else)")
    found: int = 0
    for i in range(10):
        if i == 4:
            found = 1
            break
    else:
        found = 2
    testcheck.check(found == 1, "for-else break OK (1)", "for-else break FAILED expect 1")


def test_for_else_empty_range():
    testcheck.section("Test 13: for-else (range(0) 空迭代 → 执行 else)")
    v: int = 0
    for i in range(0):
        v = 1
    else:
        v = 7
    testcheck.check(v == 7, "for-else empty range OK (7)", "for-else empty range FAILED expect 7")


def test_for_else_list():
    testcheck.section("Test 14: for-else (list[str] 迭代耗尽 → 执行 else)")
    arena: bytes = stdlib.malloc(65536)
    bd: memhub.MemBuddy* = memhub.MemBuddy(arena, 65536)
    lst: list[str] = list[str](bd)
    lst.append("a")
    lst.append("b")
    lst.append("c")
    cnt: int = 0
    ran_else: int = 0
    for v in lst:
        cnt += 1
    else:
        ran_else = 1
    testcheck.check(cnt == 3, "for-else list count OK (3)", "for-else list count FAILED expect 3")
    testcheck.check(ran_else == 1, "for-else list else OK (1)", "for-else list else FAILED expect 1")
    stdlib.free(arena)


def test_while_else_cond_false():
    testcheck.section("Test 15: while-else (条件变假 → 执行 else)")
    i: int = 0
    acc: int = 0
    while i < 3:
        acc += i
        i += 1
    else:
        acc += 100
    testcheck.check(acc == 103, "while-else cond-false OK (103)", "while-else cond-false FAILED expect 103")


def test_while_else_zero_iter():
    testcheck.section("Test 16: while-else (一次都不进循环 → 执行 else)")
    n: int = 0
    while n < 0:
        n = 99
    else:
        n = 7
    testcheck.check(n == 7, "while-else zero-iter OK (7)", "while-else zero-iter FAILED expect 7")


def test_while_else_break():
    testcheck.section("Test 17: while-else (break → 跳过 else)")
    i: int = 0
    flag: int = 0
    while i < 10:
        if i == 2:
            flag = 1
            break
        i += 1
    else:
        flag = 2
    testcheck.check(flag == 1, "while-else break OK (1)", "while-else break FAILED expect 1")


def main() -> int:
    testcheck.begin("ForAndLoopTest: 迭代器与循环测试")

    test_for_in_range()
    test_for_in_str()
    test_for_in_list()
    test_zip_list_list()
    test_zip_range_list()
    test_zip_str_str()
    test_zip_range_len_list()
    test_enumerate_list()
    test_enumerate_str()
    test_enumerate_range()
    test_for_else_no_break()
    test_for_else_break()
    test_for_else_empty_range()
    test_for_else_list()
    test_while_else_cond_false()
    test_while_else_zero_iter()
    test_while_else_break()

    return testcheck.end()