import socket
from time import sleep
import asyncio
import argparse
import os


class RPSCPI:
    def __init__(self, rp_name, rp_host):
        self.name = rp_name
        self.host = rp_host

    def connect(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.host, 5000))
        self.sock_f = self.sock.makefile()

    def close(self):
        self.sock_f.close()
        self.sock.close()

    def sendmsg(self, msg):
        self.sock.send(msg.encode() + b"\r\n")

    def recvmsg(self):
        return self.sock_f.readline().strip()

    def trigger_prep(self):
        self.sendmsg("ACQ:RST")
        self.sendmsg("ACQ:START")

    async def trigger(self):
        self.sendmsg("ACQ:TRIG NOW")
        while True:
            self.sendmsg("ACQ:TRIG:STAT?")
            if self.recvmsg() == "TD":
                break

    def get_data(self, channel):
        self.sendmsg("ACQ:SOUR{}:DATA?".format(channel))
        return self.recvmsg()[1:-1]


RP_IP_ADDRS = {
    "creotech-1": "192.168.1.104",
    "creotech-2": "192.168.1.105",
    "mlabs": "rp-f05cc9",
}


async def gather_trigger(rp_list):
    await asyncio.gather(*[rp.trigger() for rp in rp_list])


def main():
    parser = argparse.ArgumentParser(description="Data collection tool for Sayma DAC/TTL at RedPitaya")
    parser.add_argument("dir", help="output directory", type=str)
    parser.add_argument("rps", help="name(s) of the target RedPitayas where data is collected "
        "simultaneously; any of: " + " ".join(list(RP_IP_ADDRS.keys())),
        type=str, nargs='+')
    parser.add_argument("--txt", help="save data as additional human-readable text files",
        action="store_true")
    args = parser.parse_args()

    # Connect a socket to each RP
    rps = []
    for rp_name in args.rps:
        rp = RPSCPI(rp_name, RP_IP_ADDRS[rp_name])
        rp.connect()
        rps.append(rp)

    try:
        for rp in rps:
            rp.trigger_prep()

        asyncio.run(gather_trigger(rps))

        for rp in rps:
            y1_raw = rp.get_data(1)
            y2_raw = rp.get_data(2)

            with open(os.path.join(args.dir, 'rp_{}_y1_raw.bin'.format(rp.name)), 'wb') as f:
                f.write(y1_raw.encode('utf-8'))
                print("Succesfully written y1 raw data from RP {}.".format(rp.name))
            with open(os.path.join(args.dir, 'rp_{}_y2_raw.bin'.format(rp.name)), 'wb') as f:
                f.write(y2_raw.encode('utf-8'))
                print("Succesfully written y2 raw data from RP {}.".format(rp.name))

            if args.txt:
                y1 = [float(i) for i in y1_raw.split(',')]
                y2 = [float(i) for i in y2_raw.split(',')]
                with open(os.path.join(args.dir, 'rp_{}_y1_raw.txt'.format(rp.name)), 'w') as f:
                    for i in y1:
                        f.write(str(i) + '\n')
                    print("Succesfully written y1 human-readable data from {}.".format(rp.name))
                with open(os.path.join(args.dir, 'rp_{}_y2_raw.txt'.format(rp.name)), 'w') as f:
                    for i in y2:
                        f.write(str(i) + '\n')
                    print("Succesfully written y2 human-readable data from {}.".format(rp.name))

    finally:
        for rp in rps:
            rp.close()

if __name__ == "__main__":
    main()