#!/usr/bin/env python3

import os
import re
import socket
import sys
import time

from random import randrange

HOST = "127.0.0.1"  # Localhost
PORT = int(sys.argv[1])  # Port to listen on
TOTAL_CH = int(sys.argv[2])  # Number of Channels
LOG2FILE = False if len(sys.argv) < 4 else bool(int(sys.argv[3]))

import logging
logger = logging.getLogger('daq')

if LOG2FILE:
    logging.basicConfig(filename=os.path.join(os.getcwd(), 'daq_sim.log'), level=logging.INFO)


class DAQ:
    def __init__(self, total_channels):
        self.total_channels = total_channels
        self.counts = [0] * self.total_channels
        self.rates = [randrange(5) for i in range(self.total_channels)]

        self.status = 0
        self.countmode = 'time'
        self.presettime = 0
        self.presetcount = 0
        self.starttime = 0
        self.elapsed = 0
        self.monitor = 0

        self.minratechannel = 0
        self.minrates = [2] * self.total_channels

        self.gate_config = [
           #(enabled, count high/low),
            (  False,          True),
            (  False,          True),
        ]

        self.gate = [
          # high/low
            False,
            False,
        ]

    def clearCount(self, counter):
        self.counts[counter-1] = 0

    def clearCounts(self):
        self.counts = [0] * self.total_channels

    def clearTime(self):
        self.elapsed = 0
        self.starttime = time.time()

    def getStatus(self):
        return self.status

    def getCount(self, channel):
        return self.counts[channel - 1]

    def getCounts(self):
        # The sinqtest daq returns a maximum of 8
        return self.counts[0:min(len(self.counts), 8)]

    def getMonitorCount(self):
        return self.counts[self.monitor]

    def updateRates(self):
        for i in range(self.total_channels):
            self.rates[i] = randrange(5)

    def updateCounts(self):
        for i in range(self.total_channels):
            self.counts[i] += self.rates[i]

    def getRunTime(self):
        elapsed = round(time.time() - self.starttime, 3)
        self.updateRates()

        # If gating and a low rate threshold are enabled, then the threshold
        # seems to have precedence and sets the status to 5.
        # If we are waiting on the gate, then the status is just 1 as normal
        # after having started a count.

        if self.countmode == 'time':
            if elapsed < self.presettime:

                if self.minratechannel >= 0 and self.rates[self.minratechannel] < self.minrates[self.minratechannel]:
                    # adjust the starttime, so that it is as if this polling period didn't happen
                    self.starttime += elapsed - self.elapsed
                    self.status = 5
                    return self.elapsed

                self.status = 1

                for i in range(len(self.gate)):
                    # If the gating is enabled and the signal is in the active position
                    if self.gate_config[i][0] and self.gate_config[i][1] != self.gate[i]:
                        self.starttime += elapsed - self.elapsed
                        return self.elapsed

                self.updateCounts()
                self.elapsed = elapsed
            else:
                self.elapsed = self.presettime if self.presettime > 0 else self.elapsed
                self.status = 0

        elif self.countmode == 'count':
            if self.getMonitorCount() < self.presetcount:

                if self.minratechannel >= 0 and self.rates[self.minratechannel] < self.minrates[self.minratechannel]:
                    # adjust the starttime, so that it is as if this polling period didn't happen
                    self.starttime += elapsed - self.elapsed
                    self.status = 5
                    return self.elapsed

                self.status = 1

                for i in range(len(self.gate)):
                    # If the gating is enabled and the signal is in the active position
                    if self.gate_config[i][0] and self.gate_config[i][1] != self.gate[i]:
                        self.starttime += elapsed - self.elapsed
                        return self.elapsed

                self.updateCounts()
                self.elapsed = elapsed

            if self.getMonitorCount() >= self.presetcount:
                self.counts[self.monitor] = self.presetcount if self.presetcount > 0 else self.counts[self.monitor]
                self.status = 0            

        else:
            raise Exception("Invalid State")

        return self.elapsed

    def stop(self):
        self.getRunTime()
        self.status = 0
        self.presettime = 0
        self.presetcount = 0

    def startTimePreset(self, presettime):
        self.countmode = 'time'
        self.status = 1
        self.presettime = round(presettime, 3)
        self.clearTime()
        self.clearCounts()

    def startCountPreset(self, presetcount):
        self.countmode = 'count'
        self.status = 1
        self.presetcount = presetcount
        self.clearTime()
        self.clearCounts()

    def setMonitorChannel(self, channel):
        self.monitor = channel - 1

    def getMonitorChannel(self):
        return self.monitor + 1

    def getRate(self, channel):
        return float(self.rates[channel - 1])

    def getMinRateChannel(self):
        return self.minratechannel + 1

    def setMinRateChannel(self, channel):
        self.minratechannel = channel - 1

    def getMinRate(self, channel):
        return self.minrates[channel - 1]

    def setMinRate(self, channel, rate):
        self.minrates[channel - 1] = rate

    def getGateStatus(self, channel):
        return self.gate_config[channel - 1]

    def setGateStatus(self, channel, enable, highlow):
        self.gate_config[channel - 1] = (enable, highlow)

    def setGate(self, channel, highlow):
        self.gate[channel - 1] = highlow

        
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:

    s.bind((HOST, PORT))
    s.listen()
    conn, addr = s.accept()
    with conn:

        def send(data: str):
            if data:
                logger.info(f'SENDING: "{data}"')
                return conn.sendall(f'{data}\r'.encode())
            else:
                logger.info(f'SENDING: ""')
                return conn.sendall(b'\r')

        def receive():
            data = conn.recv(1024)
            if data:
                # also removes terminator
                received = data.decode('ascii').rstrip()
            else:
                received = ''
            logger.info(f'RECEIVED: "{received}"')
            return received
        
        daq = DAQ(TOTAL_CH)

        while True:

            data = receive()

            if not data: # Empty implies client disconnected
                break    # Close Server

            try:

                if data == 'RMT 1':
                    send('')

                elif data == 'ECHO 2':
                    send('DAQ System Version Simulation') # Sends some sort of info command

                elif data == 'RA':
                    send(
                        ' '.join(map(str,
                            [daq.getRunTime()] + \
                            daq.getCounts()
                        ))
                    )

                elif re.fullmatch(r'RC (\d+)', data):
                    channel = int(re.fullmatch(r'RC (\d+)', data).group(1))
                    count = daq.getCount(channel)
                    send(f'{count}')

                elif data == 'RS':
                    send(str(daq.getStatus()))

                elif data == 'S':
                    daq.stop()
                    send('')

                elif data == 'CT':
                    daq.clearTime()
                    send('')

                elif re.fullmatch(r'CC (\d+)', data):
                    counter = int(re.fullmatch(r'CC (\d+)', data).group(1))
                    daq.clearCount(counter)
                    send('')

                elif re.fullmatch(r'TP (\d+(\.\d+)?)', data):
                    presettime = float(re.fullmatch(r'TP (\d+(\.\d+)?)', data).group(1))
                    daq.startTimePreset(presettime)
                    send('')

                elif re.fullmatch(r'MP (\d+)', data):
                    counts = int(re.fullmatch(r'MP (\d+)', data).group(1))
                    daq.startCountPreset(counts)
                    send('')

                elif data == 'PC':
                    send(str(daq.getMonitorChannel()))

                elif re.fullmatch(r'PC (\d+)', data):
                    channel = int(re.fullmatch(r'PC (\d+)', data).group(1))
                    daq.setMonitorChannel(channel)
                    send('')

                elif data == 'DR':
                    send(str(daq.getMinRateChannel()))

                elif re.fullmatch(r'DR (\d+)', data):
                    channel = int(re.fullmatch(r'DR (\d+)', data).group(1))
                    daq.setMinRateChannel(channel)
                    send('')

                elif re.fullmatch(r'DL (\d+)', data):
                    channel = int(re.fullmatch(r'DL (\d+)', data).group(1))
                    send('{:.3f}'.format(daq.getMinRate(channel)))

                elif re.fullmatch(r'DL (\d+) (\d+(?:.\d+)?)', data):
                    channel, rate = re.fullmatch(r'DL (\d+) (\d+(?:.\d+)?)', data).groups()
                    daq.setMinRate(int(channel), float(rate))
                    send('')

                elif re.fullmatch(r'RR (\d+)', data):
                    channel = int(re.fullmatch(r'RR (\d+)', data).group(1))
                    send(daq.getRate(channel))

                elif re.fullmatch(r'GT (\d+)', data):
                    channel = int(re.fullmatch(r'GT (\d+)', data).group(1))
                    enabled, highlow = daq.getGateStatus(channel)
                    send(f'{int(enabled)} {int(highlow)}')

                elif re.fullmatch(r'GT (\d+) (\d+) (\d+)', data):
                    channel, enable, highlow = re.fullmatch(r'GT (\d+) (\d+) (\d+)', data).groups()
                    channel, enable, highlow = int(channel), bool(int(enable)), bool(int(highlow))
                    daq.setGateStatus(channel, enable, highlow)
                    if enable:
                        send(f'Gate {channel} enabled, counting when input = {"high" if highlow else "low"}')
                    else:
                        send(f'Gate {channel} disabled')

                # Only for the simulation
                elif re.fullmatch(r'GATE (\d+) (\d+)', data):
                    channel, highlow = re.fullmatch(r'GATE (\d+) (\d+)', data).groups()
                    channel, highlow = int(channel), bool(int(highlow))
                    daq.setGate(channel, highlow)
                    send('')

                else:
                    send('?2') # Bad command

            except Exception as e:
                logger.exception('Simulation Broke')
                send('?OV')