"""
08/21, R James, F Alder

Adapted from script by https://github.com/awmlee
"""
import time
from websocket import create_connection
import requests
import json
import keyboard

class Websocket:
    def __init__(self, socket_url=None):
        self.rp_ip = 'rp-f05a98.local' if socket_url == '' else socket_url
        self.socket_url = 'ws://' + self.rp_ip + '/wss'
        self.start_app_url = 'http://' + self.rp_ip + '/bazaar?start=streaming_manager'
        self.stop_app_url = 'http://' + self.rp_ip + '/bazaar?stop=streaming_manager'


#------------------------------------------------------------------------------
if __name__ == '__main__':
    socket_url = input('Red Pitaya IP address (press enter for default: rp-f05a98.local):')
    WS = Websocket(socket_url)

    # Start streaming server
    try:
        web = requests.get(WS.start_app_url)
        print('Connection established')
    except:
        raise OSError('Could not connect to Red Pitaya')

    # Create websocket connection
    ws = create_connection(WS.socket_url)

    # Set websocket parameters
    setbits = ({'parameters': {'SS_RESOLUTION': {'value': 2}, 'in_command': {'value': 'send_all_params'}}})
    setrate = ({'parameters': {'SS_RATE': {'value': 800}, 'in_command': {'value': 'send_all_params'}}})
    setdualchan = ({'parameters': {'SS_CHANNEL': {'value': 3},' in_command': {'value': 'send_all_params'}}})
    ws.send(json.dumps(setbits))
    ws.send(json.dumps(setrate))
    ws.send(json.dumps(setdualchan))

    startcmd = ({'parameters': {'SS_START': {'value': 1}, 'in_command': {'value': 'send_all_params'}}})
    stopcmd = ({'parameters': {'SS_START': {'value': 0}, 'in_command': {'value': 'send_all_params'}}})

    # Start streaming data
    ws.send(json.dumps(startcmd))

    print('Streaming: press q to stop')
    while True:
        if keyboard.read_key() == 'q':
            print('Ending streaming, closing connection')
            break

    # Stop streaming data and close websocket
    ws.send(json.dumps(stopcmd))
    ws.close()
