import socket

preamble_fields = {'BYT_NR': int, # data width for waveform
                   'BIT_NR': int, # number of bits per waveform point
                   'ENCDG' : str, # encoding of waveform (binary/ascii)
                   'BN_FMT': str, # binary format of waveform
                   'BYT_OR': str, # ordering of waveform data bytes (LSB/MSB)
                   'NR_PT' : int, # record length of record waveform
                   'WFID'  : str, # description string of waveform
                   'PT_FMT': str, # format of reverence waveform (Y/ENV)
                   'XINCR' : float,
                   'PT_OFF': int,
                   'XZERO' : float,
                   'XUNIT' : str,
                   'YMULT' : float,
                   'YZERO' : float,
                   'YOFF'  : float,
                   'YUNIT' : str,
                   'NR_FR' : int }

class TekScope(object):
    def __init__(self, host, port):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((host,port))

    def send(self, msg):
        self.sock.sendall(msg)

    def recv(self):
        buffer = ''

        while True:
            buffer += self.sock.recv(1024)

            if buffer.endswith('\n'):
                break

        return buffer

    def query(self, msg):
        if not msg.endswith('\n'):
            msg += '\n'

        self.send(msg)

        return self.recv()

    def get_preamble(self):
        header = self.query('header?\n')

        # turn header on so that we know preamble field names
        self.send('header 1\n')

        preamble = {}
        for s in self.query('wfmpre?\n').strip()[8:].split(';'):
            key, value = s.split(' ',1)

            preamble[key] = preamble_fields[key](value)

        self.send(header)

        return preamble