# -*- coding: utf-8 -*-
"""
完整截取微信通讯录 - 小步滚动确保不遗漏
"""
import uiautomation as auto
import time
import os
from PIL import Image


def find_wechat_window():
    """查找微信窗口"""
    wechat_window = auto.WindowControl(searchDepth=1, Name='微信')
    if wechat_window.Exists(3, 1):
        return wechat_window
    return None


def capture_contact_region(wechat_window, index, save_dir):
    """截取通讯录区域"""
    rect = wechat_window.BoundingRectangle
    
    # 通讯录列表区域
    x_offset = 70
    y_offset = 130
    width = 280
    height = rect.height() - 160
    
    screenshot_path = os.path.join(save_dir, f"region_{index:04d}.png")
    
    try:
        bitmap = wechat_window.ToBitmap(x=x_offset, y=y_offset, width=width, height=height)
        bitmap.ToFile(screenshot_path)
        return screenshot_path
    except Exception as e:
        print(f"截图失败: {e}")
        return None


def get_image_hash(image_path):
    """计算图片感知哈希"""
    try:
        img = Image.open(image_path)
        img = img.resize((32, 32), Image.Resampling.LANCZOS)
        img = img.convert('L')
        pixels = list(img.getdata())
        avg = sum(pixels) / len(pixels)
        return ''.join(['1' if p > avg else '0' for p in pixels])
    except Exception as e:
        print(f"哈希计算失败: {e}")
        return None


def images_similarity(hash1, hash2):
    """计算图片相似度"""
    if not hash1 or not hash2:
        return 0
    diff = sum(c1 != c2 for c1, c2 in zip(hash1, hash2))
    return 1 - diff / len(hash1)


def scroll_down(wechat_window, wheel_times=1):
    """向下滚动 - 小幅度滚动"""
    try:
        rect = wechat_window.BoundingRectangle
        center_x = rect.left + 200
        center_y = rect.top + 400
        
        auto.SetCursorPos(center_x, center_y)
        auto.Click(center_x, center_y)
        time.sleep(0.15)
        auto.WheelDown(wheelTimes=wheel_times)
        time.sleep(0.25)
        return True
    except Exception as e:
        print(f"滚动失败: {e}")
        return False


def main():
    print("=" * 60)
    print("完整截取微信通讯录 - 小步滚动版本")
    print("=" * 60)
    
    # 创建新的截图目录
    save_dir = r"D:\夏骥\微信研究\scroll_complete"
    
    # 清空或创建目录
    if os.path.exists(save_dir):
        import shutil
        shutil.rmtree(save_dir)
    os.makedirs(save_dir)
    
    print(f"截图保存目录: {save_dir}")
    
    # 查找微信窗口
    print("\n查找微信窗口...")
    wechat_window = find_wechat_window()
    if not wechat_window:
        print("未找到微信窗口！请确保微信已打开。")
        return
    
    print(f"找到微信窗口: {wechat_window.Name}")
    
    # 提示用户准备
    print("\n" + "!" * 60)
    print("请确保微信通讯录界面已打开，并滚动到最顶部！")
    print("按字母顺序从A开始显示联系人列表")
    print("!" * 60)
    time.sleep(3)
    
    # 开始截图
    print("\n开始截图...")
    screenshots = []
    last_hash = None
    no_change_count = 0
    max_screenshots = 2000  # 最大截图数
    consecutive_same = 3    # 连续相同截图数判定到底
    
    for i in range(max_screenshots):
        path = capture_contact_region(wechat_window, i, save_dir)
        if not path:
            print(f"截图 {i} 失败，跳过")
            continue
        
        screenshots.append(path)
        print(f"  截图 {i + 1}: {os.path.basename(path)}")
        
        # 检测是否到底（图片相似度）
        current_hash = get_image_hash(path)
        if last_hash:
            sim = images_similarity(last_hash, current_hash)
            if sim > 0.98:  # 非常相似
                no_change_count += 1
                print(f"    -> 相似度 {sim:.2%}，连续相同 {no_change_count}")
                if no_change_count >= consecutive_same:
                    print(f"\n检测到底部！连续 {consecutive_same} 张截图相同")
                    break
            else:
                no_change_count = 0
        last_hash = current_hash
        
        # 小幅度滚动（每次滚动1格）
        scroll_down(wechat_window, wheel_times=1)
    
    print("\n" + "=" * 60)
    print(f"截图完成！共 {len(screenshots)} 张")
    print(f"保存目录: {save_dir}")
    print("=" * 60)
    
    # 提示下一步
    print("\n下一步：运行 batch_ocr_complete.py 进行OCR识别")


if __name__ == '__main__':
    main()