import subprocess
import os
import time


def take_screenshot(device_ip, device_port, save_path='screenshots'):
    # 构建连接命令
    connect_command = f'adb connect {device_ip}:{device_port}'
    subprocess.run(connect_command, shell=True, check=True)

    # 构建设备特定的截图命令
    timestamp = time.strftime('%Y%m%d_%H%M%S')
    screenshot_name = f'screenshot_{timestamp}.png'
    screenshot_path = os.path.join(save_path, screenshot_name)

    # 检查保存路径是否存在，如果不存在则创建
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # 执行截图命令
    screenshot_command = f'adb -s {device_ip}:{device_port} shell screencap -p /sdcard/{screenshot_name}'
    subprocess.run(screenshot_command, shell=True, check=True)

    # 从设备拉取截图到本地指定路径
    pull_command = f'adb -s {device_ip}:{device_port} pull /sdcard/{screenshot_name} {screenshot_path}'
    subprocess.run(pull_command, shell=True, check=True)

    # 删除设备上的截图文件
    remove_command = f'adb -s {device_ip}:{device_port} shell rm /sdcard/{screenshot_name}'
    subprocess.run(remove_command, shell=True, check=True)

    print(f'Screenshot saved as {screenshot_path}')


if __name__ == '__main__':
    device_ip = '192.168.1.100'  # 替换为你设备的IP地址
    device_port = '5555'  # 替换为你设备的端口号
    take_screenshot(device_ip, device_port)