import stdio
import stdlib
import t, c
import testcheck


# ============================================================
# 运算符重载测试
#
# 验证: 当 lhs 是结构体指针时，二元/比较运算触发 dunder 方法调用
# 支持: __add__/__sub__/__mul__/__div__/__mod__/__and__/__or__/__xor__
#       __lshift__/__rshift__/__floordiv__ (BinOp)
#       __eq__/__ne__/__lt__/__le__/__gt__/__ge__ (Compare)
#
# 三种测试场景:
#   1. Counter/BitBox: dunder 返回 t.CInt（值类型变量触发重载）
#   2. Vec2: dunder 返回 Vec2|t.CPtr（堆分配结构体，指针类型变量触发重载）
#   3. 原生整数运算回退
# ============================================================


# ============================================================
# Counter: 演示 __add__/__mul__/__sub__ + 比较重载
# dunder 返回 t.CInt
# ============================================================
class Counter:
    val: t.CInt

    def __init__(self, v: t.CInt):
        self.val = v

    def __add__(self, n: t.CInt) -> t.CInt:
        return self.val + n

    def __sub__(self, n: t.CInt) -> t.CInt:
        return self.val - n

    def __mul__(self, n: t.CInt) -> t.CInt:
        return self.val * n

    def __div__(self, n: t.CInt) -> t.CInt:
        return self.val / n

    def __mod__(self, n: t.CInt) -> t.CInt:
        return self.val % n

    def __eq__(self, n: t.CInt) -> t.CInt:
        if self.val == n:
            return 1
        return 0

    def __ne__(self, n: t.CInt) -> t.CInt:
        if self.val != n:
            return 1
        return 0

    def __lt__(self, n: t.CInt) -> t.CInt:
        if self.val < n:
            return 1
        return 0

    def __le__(self, n: t.CInt) -> t.CInt:
        if self.val <= n:
            return 1
        return 0

    def __gt__(self, n: t.CInt) -> t.CInt:
        if self.val > n:
            return 1
        return 0

    def __ge__(self, n: t.CInt) -> t.CInt:
        if self.val >= n:
            return 1
        return 0


# ============================================================
# BitBox: 演示位运算重载
# dunder 返回 t.CInt
# ============================================================
class BitBox:
    flags: t.CInt

    def __init__(self, v: t.CInt):
        self.flags = v

    def __and__(self, mask: t.CInt) -> t.CInt:
        return self.flags & mask

    def __or__(self, mask: t.CInt) -> t.CInt:
        return self.flags | mask

    def __xor__(self, mask: t.CInt) -> t.CInt:
        return self.flags ^ mask

    def __lshift__(self, n: t.CInt) -> t.CInt:
        return self.flags << n

    def __rshift__(self, n: t.CInt) -> t.CInt:
        return self.flags >> n


# ============================================================
# Vec2: 演示结构体运算重载
#
# __new__ + stdlib.malloc 在堆上分配内存，使 Vec2|t.CPtr 变量
# 能正确工作（dunder 方法返回堆上新对象）。
#
# __new__ 返回 malloc 的堆指针作为结构体存储空间，
# 后续 __init__ 在堆指针上初始化字段。
# ============================================================
class Vec2:
    x: t.CInt
    y: t.CInt

    def __new__(self, x0: t.CInt, y0: t.CInt):
        r: Vec2 | t.CPtr = stdlib.malloc(Vec2.__sizeof__())
        return r

    def __init__(self, x0: t.CInt, y0: t.CInt):
        self.x = x0
        self.y = y0

    def __add__(self, other: Vec2 | t.CPtr) -> Vec2 | t.CPtr:
        r: Vec2 | t.CPtr = Vec2(0, 0)
        r.x = self.x + other.x
        r.y = self.y + other.y
        return r

    def __sub__(self, other: Vec2 | t.CPtr) -> Vec2 | t.CPtr:
        r: Vec2 | t.CPtr = Vec2(0, 0)
        r.x = self.x - other.x
        r.y = self.y - other.y
        return r

    def __eq__(self, other: Vec2 | t.CPtr) -> t.CInt:
        if self.x == other.x:
            if self.y == other.y:
                return 1
        return 0

    def __ne__(self, other: Vec2 | t.CPtr) -> t.CInt:
        if self.x != other.x:
            return 1
        if self.y != other.y:
            return 1
        return 0


# ============================================================
# 主测试入口
# ============================================================
def opovl_test() -> int:
    testcheck.begin("Operator Overload Tests")

    # ============================================================
    # 测试 1: Counter + int 触发 __add__
    # ============================================================
    testcheck.section("Counter __add__")
    cnt: Counter = Counter(10)
    sum_val: int = cnt + 5
    testcheck.check(sum_val == 15, "cnt(10)+5=15", "cnt+5 expected 15 got ??")

    # ============================================================
    # 测试 2: Counter - int / * int / / int / % int
    # ============================================================
    testcheck.section("Counter __sub__/__mul__/__div__/__mod__")
    sub_val: int = cnt - 3
    mul_val: int = cnt * 3
    div_val: int = cnt / 3
    mod_val: int = cnt % 3
    testcheck.check(sub_val == 7, "cnt-3=7", "cnt-3 expected 7")
    testcheck.check(mul_val == 30, "cnt*3=30", "cnt*3 expected 30")
    testcheck.check(div_val == 3, "cnt/3=3", "cnt/3 expected 3")
    testcheck.check(mod_val == 1, "cnt%3=1", "cnt%3 expected 1")

    # ============================================================
    # 测试 3: Counter 比较重载 (==, !=, <, <=, >, >=)
    # ============================================================
    testcheck.section("Counter compare overloads")
    eq_10: int = cnt == 10
    ne_10: int = cnt != 10
    lt_20: int = cnt < 20
    le_10: int = cnt <= 10
    gt_5: int = cnt > 5
    ge_10: int = cnt >= 10
    testcheck.check(eq_10 == 1, "cnt==10:1", "cnt==10 expected 1")
    testcheck.check(ne_10 == 0, "cnt!=10:0", "cnt!=10 expected 0")
    testcheck.check(lt_20 == 1, "cnt<20:1", "cnt<20 expected 1")
    testcheck.check(le_10 == 1, "cnt<=10:1", "cnt<=10 expected 1")
    testcheck.check(gt_5 == 1, "cnt>5:1", "cnt>5 expected 1")
    testcheck.check(ge_10 == 1, "cnt>=10:1", "cnt>=10 expected 1")

    # ============================================================
    # 测试 4: BitBox 位运算重载
    # ============================================================
    testcheck.section("BitBox bitwise overloads")
    bb: BitBox = BitBox(0xFF)
    and_val: int = bb & 0x0F
    or_val: int = bb | 0x100
    xor_val: int = bb ^ 0x55
    lsh_val: int = bb << 4
    rsh_val: int = bb >> 4
    testcheck.check(and_val == 0x0F, "bb&0x0F=15", "bb&0x0F expected 15")
    testcheck.check(or_val == 0x1FF, "bb|0x100=511", "bb|0x100 expected 511")
    testcheck.check(xor_val == 0xAA, "bb^0x55=170", "bb^0x55 expected 170")
    testcheck.check(lsh_val == 0xFF0, "bb<<4=4080", "bb<<4 expected 4080")
    testcheck.check(rsh_val == 0x0F, "bb>>4=15", "bb>>4 expected 15")

    # ============================================================
    # 测试 5: 原生整数运算不受影响（回退路径）
    # ============================================================
    testcheck.section("Native int binop")
    x: int = 7 + 8
    y: int = 10 - 3
    z: int = 4 * 5
    testcheck.check(x == 15, "7+8=15", "7+8 expected 15")
    testcheck.check(y == 7, "10-3=7", "10-3 expected 7")
    testcheck.check(z == 20, "4*5=20", "4*5 expected 20")

    # ============================================================
    # 测试 6: 原生整数比较不受影响
    # ============================================================
    testcheck.section("Native int compare")
    testcheck.check(3 < 5, "3<5", "3<5 broken")
    testcheck.check(10 == 10, "10==10", "10==10 broken")

    # ============================================================
    # 测试 7: Vec2 结构体运算重载
    # ============================================================
    testcheck.section("Vec2 struct overloads")
    a: Vec2 | t.CPtr = Vec2(3, 4)
    b: Vec2 | t.CPtr = Vec2(1, 2)
    c: Vec2 | t.CPtr = a + b
    d: Vec2 | t.CPtr = a - b
    testcheck.check(c.x == 4, "a+b x=4", "a+b x expected 4")
    testcheck.check(c.y == 6, "a+b y=6", "a+b y expected 6")
    testcheck.check(d.x == 2, "a-b x=2", "a-b x expected 2")
    testcheck.check(d.y == 2, "a-b y=2", "a-b y expected 2")

    # ============================================================
    # 测试 8: Vec2 比较重载
    # ============================================================
    testcheck.section("Vec2 compare overloads")
    p: Vec2 | t.CPtr = Vec2(5, 5)
    q: Vec2 | t.CPtr = Vec2(5, 5)
    eq_v: int = p == q
    ne_v: int = p != q
    testcheck.check(eq_v == 1, "(5,5)==(5,5):1", "p==q expected 1")
    testcheck.check(ne_v == 0, "(5,5)!=(5,5):0", "p!=q expected 0")

    return testcheck.end()