import ctypes, time, sys, os
from PIL import ImageGrab
import win32gui, win32process, psutil

sys.stdout.reconfigure(encoding='utf-8')
log_file = r"C:\Projects\JX\logs\elevated_test_run.log"

def log(msg):
    ts = time.strftime("[%Y-%m-%d %H:%M:%S]")
    line = f"{ts} {msg}\n"
    print(line, end='')
    with open(log_file, "a", encoding="utf-8") as f:
        f.write(line)

with open(log_file, "w", encoding="utf-8") as f:
    f.write("=== ELEVATED TEST RUN START ===\n")

user32 = ctypes.windll.user32
shell32 = ctypes.windll.shell32
kernel32 = ctypes.windll.kernel32

is_admin = bool(shell32.IsUserAnAdmin())
log(f"Process PID: {os.getpid()}, IsUserAnAdmin: {is_admin}")
if not is_admin:
    log("FATAL: Not running as Administrator! Exiting.")
    sys.exit(1)

hdesk = user32.OpenInputDesktop(0, False, 0x01FF)
if hdesk:
    user32.SetThreadDesktop(hdesk)

# Check for existing game window
existing_hwnds = []
def enum_cb(h, _):
    if win32gui.GetClassName(h) == "Sword3 Class":
        existing_hwnds.append(h)
win32gui.EnumWindows(enum_cb, None)
log(f"Existing Sword3 windows: {[hex(h) for h in existing_hwnds]}")

target_hwnd = None
if existing_hwnds:
    target_hwnd = existing_hwnds[0]
    log(f"Using existing window: {hex(target_hwnd)}")
else:
    log("Launching Game.exe with runas...")
    res = shell32.ShellExecuteW(None, "runas", r"C:\JXVCTK\JxVoCongTruyenKy\Game.exe", None, r"C:\JXVCTK\JxVoCongTruyenKy", 1)
    log(f"ShellExecuteW returned: {res}")
    
    start_t = time.time()
    while time.time() - start_t < 25:
        time.sleep(1.0)
        current = []
        def cb2(h, _):
            if win32gui.GetClassName(h) == "Sword3 Class":
                current.append(h)
        win32gui.EnumWindows(cb2, None)
        diff = set(current) - set(existing_hwnds)
        if diff:
            target_hwnd = list(diff)[0]
            log(f"Found new Sword3 window: {hex(target_hwnd)} after {time.time()-start_t:.1f}s")
            break

if not target_hwnd or not win32gui.IsWindow(target_hwnd):
    log("FATAL: Game window not found!")
    sys.exit(1)

def capture(name):
    rect = win32gui.GetWindowRect(target_hwnd)
    img = ImageGrab.grab(bbox=rect, all_screens=True)
    path = rf"c:\Projects\JX\logs\{name}.png"
    img.save(path)
    log(f"Saved screenshot: {path}")

# Wait for DirectX load
log("Waiting 4.0s for DirectX render...")
time.sleep(4.0)
capture("test_0_initial")

# Bring to foreground & focus
user32.SetForegroundWindow(target_hwnd)
user32.PostMessageW(target_hwnd, 0x0006, 1, 0) # WA_ACTIVE
user32.PostMessageW(target_hwnd, 0x0007, 0, 0) # WM_SETFOCUS
time.sleep(0.3)

# Click "Bắt đầu trò chơi"
log("Clicking 'Bắt đầu trò chơi' at (400, 385)...")
lp = ((385 & 0xFFFF) << 16) | (400 & 0xFFFF)

kernel32.SetLastError(0)
r1 = user32.PostMessageW(target_hwnd, 0x0200, 0, lp)
err1 = kernel32.GetLastError()
time.sleep(0.04)

r2 = user32.PostMessageW(target_hwnd, 0x0201, 1, lp)
err2 = kernel32.GetLastError()
time.sleep(0.08)

r3 = user32.PostMessageW(target_hwnd, 0x0202, 0, lp)
err3 = kernel32.GetLastError()
time.sleep(0.1)

log(f"Click PostMessages: r1={r1}(err={err1}), r2={r2}(err={err2}), r3={r3}(err={err3})")

# Send Enter
sc = user32.MapVirtualKeyW(0x0D, 0)
lp_down = 1 | (sc << 16)
lp_up = 1 | (sc << 16) | (3 << 30)

r4 = user32.PostMessageW(target_hwnd, 0x0100, 0x0D, lp_down)
err4 = kernel32.GetLastError()
time.sleep(0.05)

r5 = user32.PostMessageW(target_hwnd, 0x0101, 0x0D, lp_up)
err5 = kernel32.GetLastError()
time.sleep(2.0)

log(f"Enter PostMessages: r4={r4}(err={err4}), r5={r5}(err={err5})")

capture("test_1_after_title_advance")
log("=== ELEVATED TEST RUN COMPLETED ===")