import pyvisa
from datetime import datetime
import os

# ==========================================
SCOPE_IP = '192.168.1.100'  # <--- 示波器IP修改这里！！！
# ==========================================

def capture_and_save():
    """连接示波器，截图并保存到当前文件夹"""
    
    print("1. 正在连接示波器...")
    try:
        # 创建VISA资源管理器
        rm = pyvisa.ResourceManager()
        
        # 构造示波器的网络地址 (端口5025是SIGLENT的标准SCPI端口)
        resource_string = f'TCPIP0::{SCOPE_IP}::5025::SOCKET'
        
        # 打开连接
        scope = rm.open_resource(resource_string)
        scope.timeout = 10000  # 10秒超时
        scope.read_termination = '\n'
        scope.write_termination = '\n'
        
        # 发送 *IDN? 命令确认连接成功
        idn = scope.query('*IDN?')
        print(f"   ✅ 连接成功！示波器型号: {idn}")
        
        # ==========================================
        # 2. 核心操作：让示波器发送屏幕截图
        # ==========================================
        print("2. 正在获取屏幕截图...")
        
        # 向示波器发送“打印屏幕为PNG”的命令
        scope.write(':PRINt? PNG')
        
        # 读取示波器返回的二进制图片数据
        img_data = scope.read_raw()
        print(f"   ✅ 截图数据已接收，大小: {len(img_data)} 字节")
        
        # ==========================================
        # 3. 保存为PNG文件
        # ==========================================
        # 生成文件名：以当前时间命名，避免覆盖
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = f'screenshot_{timestamp}.png'
        
        # 保存到当前目录
        with open(filename, 'wb') as f:
            # 尝试直接保存，如果打不开，尝试改为 f.write(img_data[11:])
            f.write(img_data)
        
        print(f"   ✅ 截图已保存为: {os.path.abspath(filename)}")
        
        # 4. 关闭连接
        scope.close()
        rm.close()
        print("3. 连接已关闭。全部完成！")
        
    except Exception as e:
        print(f"❌ 出错了！错误信息: {e}")
        print("\n请检查：")
        print("1. 示波器是否开机？")
        print("2. 网线是否插好？")
        print(f"3. IP地址 '{SCOPE_IP}' 是否和示波器屏幕显示的一致？")
        print("4. 电脑和示波器是否在同一个局域网？")

# 如果直接运行这个脚本，就执行截图函数
if __name__ == '__main__':
    capture_and_save()