import ctypes
import sys
import time
from ctypes import wintypes
import win32gui
from PIL import ImageGrab

sys.stdout.reconfigure(encoding='utf-8')
user32 = ctypes.windll.user32
hdesk = user32.OpenInputDesktop(0, False, 0x01FF)
if hdesk:
    user32.SetThreadDesktop(hdesk)

hwnds = []
def enum_cb(hwnd, _):
    txt_len = user32.GetWindowTextLengthW(hwnd) + 1
    txt_buf = ctypes.create_unicode_buffer(txt_len)
    user32.GetWindowTextW(hwnd, txt_buf, txt_len)
    if "JXHelper" in txt_buf.value and user32.IsWindowVisible(hwnd):
        hwnds.append(hwnd)
    return True

WNDENUMPROC = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HWND, wintypes.LPARAM)
user32.EnumWindows(WNDENUMPROC(enum_cb), 0)

if not hwnds:
    print("JXHelper window not found!")
    sys.exit(1)

main_hwnd = hwnds[0]
print(f"Main HWND: {hex(main_hwnd)}")

children = []
def child_cb(h, _):
    cls_buf = ctypes.create_unicode_buffer(256)
    user32.GetClassNameW(h, cls_buf, 256)
    r = win32gui.GetWindowRect(h)
    mr = win32gui.GetWindowRect(main_hwnd)
    rel = (r[0] - mr[0], r[1] - mr[1], r[2] - r[0], r[3] - r[1])
    children.append((h, cls_buf.value, rel))
    return True

user32.EnumChildWindows(main_hwnd, WNDENUMPROC(child_cb), 0)

lvw = None
cbo_sect = None
for h, cls, rel in children:
    if "SysListView32" in cls:
        lvw = h
        print(f"ListView: {hex(h)}, rel: {rel}")
    elif "ComboBox" in cls and rel[2] == 130 and rel[3] == 21:
        cbo_sect = h
        print(f"Found cboSect: {hex(h)}, rel: {rel}")

WM_LBUTTONDOWN = 0x0201
WM_LBUTTONUP = 0x0202
MK_LBUTTON = 0x0001
WM_KEYDOWN = 0x0100
WM_KEYUP = 0x0101
VK_DOWN = 0x28
VK_RETURN = 0x0D

def make_lparam(x, y):
    return (y << 16) | (x & 0xFFFF)

# 1. Click Item 0 (ConCuaBoMe)
print("\n--- STEP 1: Select ConCuaBoMe (Item 0) ---")
user32.PostMessageW(lvw, WM_LBUTTONDOWN, MK_LBUTTON, make_lparam(50, 22))
user32.PostMessageW(lvw, WM_LBUTTONUP, 0, make_lparam(50, 22))
time.sleep(0.8)

rect = win32gui.GetWindowRect(main_hwnd)
img1 = ImageGrab.grab(bbox=rect)
img1.save(r"C:\Users\Admin\.gemini\antigravity\brain\5f637e30-06d3-43e8-a064-a259fea7be1a\test_concuabome_init.png")

# 2. Change Sect to Võ Đang on ConCuaBoMe
if cbo_sect:
    print("\n--- STEP 2: Change Sect on cbo_sect using click & arrow keys ---")
    # Click to focus cbo_sect
    user32.PostMessageW(cbo_sect, WM_LBUTTONDOWN, MK_LBUTTON, make_lparam(50, 10))
    user32.PostMessageW(cbo_sect, WM_LBUTTONUP, 0, make_lparam(50, 10))
    time.sleep(0.3)
    # Send Down Arrow to go from Cái Bang (idx 0) to Võ Đang (idx 1)
    user32.PostMessageW(cbo_sect, WM_KEYDOWN, VK_DOWN, 0)
    user32.PostMessageW(cbo_sect, WM_KEYUP, VK_DOWN, 0)
    time.sleep(0.3)
    user32.PostMessageW(cbo_sect, WM_KEYDOWN, VK_RETURN, 0)
    user32.PostMessageW(cbo_sect, WM_KEYUP, VK_RETURN, 0)
    time.sleep(0.8)

img2 = ImageGrab.grab(bbox=rect)
img2.save(r"C:\Users\Admin\.gemini\antigravity\brain\5f637e30-06d3-43e8-a064-a259fea7be1a\test_concuabome_changed_vodang.png")
print("Saved test_concuabome_changed_vodang.png")

# 3. Switch to BoCuaCacCon (Item 1)
print("\n--- STEP 3: Switch to BoCuaCacCon (Item 1) ---")
user32.PostMessageW(lvw, WM_LBUTTONDOWN, MK_LBUTTON, make_lparam(50, 42))
user32.PostMessageW(lvw, WM_LBUTTONUP, 0, make_lparam(50, 42))
time.sleep(0.8)

img3 = ImageGrab.grab(bbox=rect)
img3.save(r"C:\Users\Admin\.gemini\antigravity\brain\5f637e30-06d3-43e8-a064-a259fea7be1a\test_bocuacaccon_isolated_caibang.png")
print("Saved test_bocuacaccon_isolated_caibang.png")

# 4. Switch back to ConCuaBoMe (Item 0)
print("\n--- STEP 4: Switch back to ConCuaBoMe (Item 0) ---")
user32.PostMessageW(lvw, WM_LBUTTONDOWN, MK_LBUTTON, make_lparam(50, 22))
user32.PostMessageW(lvw, WM_LBUTTONUP, 0, make_lparam(50, 22))
time.sleep(0.8)

img4 = ImageGrab.grab(bbox=rect)
img4.save(r"C:\Users\Admin\.gemini\antigravity\brain\5f637e30-06d3-43e8-a064-a259fea7be1a\test_concuabome_retained_vodang.png")
print("Saved test_concuabome_retained_vodang.png")