#!/usr/bin/env python3

import xml.etree.ElementTree as ET
import socket
import argparse, sys, re, time, os


MEDUSASTATUSPORT = 21200
MEDUSACOMMANDPORT = 21100
MEDUSASERVER = 'medusa-srv0.atnf.csiro.au'

parser = argparse.ArgumentParser()
parser.add_argument('-d', '--debug', '-debug', help="Debug mode - do not send config to Medusa", action="store_true")
parser.add_argument('-spip', '--spip', help="Read SPIP XML from file, not socket. Enables debug mode")
parser.add_argument('-start', '--start', help="Start Medusa streaming", action="store_true")
parser.add_argument('-stop', '--stop', help="Stop Medusa streaming", action="store_true")
args = parser.parse_args()

# From stack overflow
#     https://stackoverflow.com/questions/28813876/how-do-i-get-pythons-elementtree-to-pretty-print-to-an-xml-file

def _pretty_print(current, parent=None, index=-1, depth=0):
    for i, node in enumerate(current):
        _pretty_print(node, current, i, depth + 1)
    if parent is not None:
        if index == 0:
            parent.text = '\n' + ('  ' * depth)
        else:
            parent[index - 1].tail = '\n' + ('  ' * depth)
        if index == len(parent) - 1:
            current.tail = '\n' + ('  ' * (depth - 1))

def getElem(parent, elem):
    child = parent.find(elem)
    if child is None:
        sys.exit("Error: Could not find <{}>".format(elem))
    return child.text

debug = args.debug

if args.spip is not None:
    # Read Medusa setup from local file (for debugging)
    debug = True
    tree = ET.parse(args.spip)
    spip_setup = tree.getroot()
    
else:
    # Create socket to Medusa server for status
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((MEDUSASERVER, MEDUSASTATUSPORT))

    # Form XML with configuration request
    command_root = ET.Element("tcs_state_request")
    ET.SubElement(command_root, "requestor").text = "search page"
    ET.SubElement(command_root, "type").text = "state"
    command_tree = ET.ElementTree(command_root)
    command = ET.tostring(command_root, encoding="ISO-8859-1")
    command += b'\r\n'

    # Send XML to server and get response
    sock.send(command)
    time.sleep(1)
    resp = sock.recv(32000)
    sock.close()

    spip_setup = ET.XML(resp)

    if debug:
        _pretty_print(spip_setup)
        ET.dump(spip_setup)
    
beam = spip_setup.find("beam")
medusa_state = beam.attrib['state']

instConf = spip_setup.find("instrument_configuration")
nStream = int(instConf.attrib['nstream'])

#if not debug and medusa_state != 'Idle':
#    print("Warning: Medusa system currently \"{}\". Will not try and configure system".format(medusa_state))
#    sys.exit(1)

    
def SPIP_control(start):
    root = ET.Element("obs_cmd")

    command = ET.SubElement(root, "command")
    if (start):
        command.text = "start"
    else:
        command.text = "stop"

    beam_config = ET.SubElement(root, "beam_configuration")
    ET.SubElement(beam_config, "nbeam", key="NBEAM").text = "1"
    ET.SubElement(beam_config, "beam_state_0", key="BEAM_STATE_0", name="1").text = "1"

    stream_config = ET.SubElement(root, "stream_configuration")
    ET.SubElement(stream_config, "nstream", key="NSTREAM").text = str(nStream)
    for i in range(nStream):
        ET.SubElement(stream_config, "active", key="STREAM_ACTIVE").text = "1"

    obs_parm = ET.SubElement(root, "observation_parameters")
    if start:
        ET.SubElement(obs_parm, "utc_start", key="UTC_START").text = "None"
    else:
        ET.SubElement(obs_parm, "utc_stop", key="UTC_STOP").text = "None"

    return(root)

if args.stop:
    if medusa_state != 'Recording':
        print("Medusa currently in \'{}\' state. Cannot stop".format(medusa_state))
        sys.exit(1)
    control_xml = SPIP_control(False)

elif args.start:
    if medusa_state != 'Configured':
        print("Medusa currently in \'{}\' state. Needs to be \'Configured\' to start".format(medusa_state))
        sys.exit(1)
    control_xml = SPIP_control(True)

else:
    print("No control given. Medusa currently \'{}\'".format(medusa_state))
    sys.exit(1)
    
if (debug):
    _pretty_print(control_xml)
    ET.dump(control_xml)
else:
    # Send control XML to server
    print("Sending control XML to Medusa")

    # Create socket to Medusa server
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((MEDUSASERVER, MEDUSACOMMANDPORT))
    
    control = ET.tostring(control_xml, encoding="ISO-8859-1")
    control += b'\r\n'
    
    sock.send(control)
    time.sleep(2)

    print("Waiting for response")
    resp = sock.recv(32000)
    sock.close()

    spip_reply= ET.XML(resp)
    _pretty_print(spip_reply)
    print(ET.tostring(spip_reply))