#!/usr/bin/env python

import argparse
import sys, time
import pyvisa as visa
from pyvisa.constants import StatusCode

from PIL import Image
from io import BytesIO

def screendump(device, filename):
    rm = visa.ResourceManager('@py')

    # Siglent device
    scope = rm.open_resource('TCPIP::%s' % device)
    scope.timeout = 1000 # ms
    scope.read_trermination = None

    print(scope.query('*IDN?'))
    scope.write('SCDP')
    print(scope.last_status)

    if scope.last_status == StatusCode.success:
        time.sleep(1)
        fp = BytesIO()
        count = 0
        try:
            data = scope.read_raw(20000000)
            fp.write(data)
            count = len(data)
        except:
            print("timeout")
            pass
        print("Read {} bytes.".format(count))
        image = Image.open(fp)
        image.save(filename)
    scope.close()
    rm.close()

if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        description='Grab a screenshot from a Siglent LXI device.')
    parser.add_argument('--output', '-o', dest='filename', required=True,
        help='the output filename')
    parser.add_argument('device', metavar='DEVICE', nargs=1,
        help='the ip address or hostname of the DSO')
    
    args = parser.parse_args()
    screendump(args.device[0], args.filename)