"""Save the pairing QR displayed by the already-running WhatsApp Desktop app.

No browser, debug port, process injection, restart, or logout. Requires Pillow
and zxing-cpp. The official Desktop app must be displaying its linking QR.
"""
import argparse
import ctypes as C
from ctypes import wintypes as W
import hashlib
import json
from pathlib import Path
import subprocess
import time


user = C.WinDLL('user32', use_last_error=True)
gdi = C.WinDLL('gdi32', use_last_error=True)
user.GetDC.argtypes = [W.HWND]; user.GetDC.restype = W.HDC
user.ReleaseDC.argtypes = [W.HWND, W.HDC]
user.GetWindowThreadProcessId.argtypes = [W.HWND, C.POINTER(W.DWORD)]
user.GetWindowRect.argtypes = [W.HWND, C.POINTER(W.RECT)]
user.PrintWindow.argtypes = [W.HWND, W.HDC, W.UINT]
user.IsWindowVisible.argtypes = [W.HWND]
gdi.CreateCompatibleDC.argtypes = [W.HDC]; gdi.CreateCompatibleDC.restype = W.HDC
gdi.CreateCompatibleBitmap.argtypes = [W.HDC, C.c_int, C.c_int]; gdi.CreateCompatibleBitmap.restype = W.HBITMAP
gdi.SelectObject.argtypes = [W.HDC, W.HGDIOBJ]; gdi.SelectObject.restype = W.HGDIOBJ
gdi.DeleteObject.argtypes = [W.HGDIOBJ]
gdi.DeleteDC.argtypes = [W.HDC]

class Header(C.Structure):
    _fields_ = [('size', W.DWORD), ('width', W.LONG), ('height', W.LONG),
                ('planes', W.WORD), ('bits', W.WORD), ('compression', W.DWORD),
                ('image_size', W.DWORD), ('xppm', W.LONG), ('yppm', W.LONG),
                ('used', W.DWORD), ('important', W.DWORD)]

class Info(C.Structure):
    _fields_ = [('header', Header), ('colors', W.DWORD * 3)]

gdi.GetDIBits.argtypes = [W.HDC, W.HBITMAP, W.UINT, W.UINT, C.c_void_p, C.POINTER(Info), W.UINT]

def desktop_windows():
    query = "@(Get-CimInstance Win32_Process -Filter \"Name='WhatsApp.Root.exe'\" | Select-Object -ExpandProperty ProcessId) | ConvertTo-Json -Compress"
    result = subprocess.run(['powershell.exe', '-NoProfile', '-NonInteractive', '-Command', query], capture_output=True, text=True, check=True)
    data = json.loads(result.stdout) if result.stdout.strip() else []
    pids = set(data if isinstance(data, list) else [data])
    windows = []
    callback_type = C.WINFUNCTYPE(W.BOOL, W.HWND, W.LPARAM)
    @callback_type
    def callback(hwnd, _):
        pid = W.DWORD()
        user.GetWindowThreadProcessId(hwnd, C.byref(pid))
        if pid.value in pids and user.IsWindowVisible(hwnd):
            windows.append((hwnd, pid.value))
        return True
    user.EnumWindows.argtypes = [callback_type, W.LPARAM]
    user.EnumWindows(callback, 0)
    return windows

def capture(hwnd):
    from PIL import Image
    rect = W.RECT()
    if not user.GetWindowRect(hwnd, C.byref(rect)):
        return None
    width, height = rect.right - rect.left, rect.bottom - rect.top
    if not (100 <= width <= 10000 and 100 <= height <= 10000):
        return None
    screen = user.GetDC(hwnd)
    dc = gdi.CreateCompatibleDC(screen)
    bitmap = gdi.CreateCompatibleBitmap(screen, width, height)
    old = gdi.SelectObject(dc, bitmap)
    try:
        if not user.PrintWindow(hwnd, dc, 2):
            return None
        gdi.SelectObject(dc, old)
        info = Info(Header(C.sizeof(Header), width, -height, 1, 32, 0, 0, 0, 0, 0, 0))
        pixels = C.create_string_buffer(width * height * 4)
        if gdi.GetDIBits(dc, bitmap, 0, height, pixels, C.byref(info), 0) != height:
            return None
        return Image.frombytes('RGB', (width, height), pixels.raw, 'raw', 'BGRX')
    finally:
        gdi.SelectObject(dc, old)
        gdi.DeleteObject(bitmap)
        gdi.DeleteDC(dc)
        user.ReleaseDC(hwnd, screen)

def pairing_payload(text):
    import base64
    parts = text.split(',')
    if len(parts) not in (5, 6):
        return False
    try:
        return bool(parts[0]) and all(len(base64.b64decode(p, validate=True)) == 32 for p in parts[1:4]) and parts[4].isdigit()
    except ValueError:
        return False

def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=Path.cwd() / 'whatsapp-desktop-qr.png')
    parser.add_argument('--wait-seconds', type=int, default=0)
    args = parser.parse_args(argv)
    import zxingcpp
    if not 0 <= args.wait_seconds <= 3600:
        parser.error('--wait-seconds must be 0..3600')
    deadline = time.monotonic() + args.wait_seconds
    while True:
        windows = desktop_windows()
        for hwnd, pid in windows:
            image = capture(hwnd)
            if image is None:
                continue
            for code in zxingcpp.read_barcodes(image):
                if code.format != zxingcpp.BarcodeFormat.QRCode or not pairing_payload(code.text):
                    continue
                points = [code.position.top_left, code.position.top_right, code.position.bottom_left, code.position.bottom_right]
                margin = 25
                box = (max(0, min(p.x for p in points)-margin), max(0, min(p.y for p in points)-margin),
                       min(image.width, max(p.x for p in points)+margin), min(image.height, max(p.y for p in points)+margin))
                cropped = image.crop(box)
                if not any(c.text == code.text for c in zxingcpp.read_barcodes(cropped)):
                    continue
                args.output.parent.mkdir(parents=True, exist_ok=True)
                cropped.save(args.output, format='PNG')
                print(json.dumps({'status': 'saved', 'source': 'running WhatsApp Desktop window', 'pid': pid,
                                  'file': str(args.output.resolve()), 'payload_sha256': hashlib.sha256(code.text.encode()).hexdigest()}))
                return 0
        if time.monotonic() >= deadline:
            print(json.dumps({'status': 'no_active_desktop_qr', 'windows_found': len(windows),
                              'note': 'Desktop must display its linking QR. A linked account has no login QR; GPU rendering or a minimized window can also prevent capture. No QR file was created or overwritten.'}))
            return 2
        time.sleep(2)

if __name__ == '__main__':
    raise SystemExit(main())