#!/usr/bin/env python3

import xml.etree.ElementTree as ET
import socket
import time, sys

# 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))

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",)

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect(("medusa-srv0.atnf.csiro.au", 21200))

s.send(command)
time.sleep(2)
resp = s.recv(32000)
s.close()

root = ET.XML(resp)
_pretty_print(root)
tree = ET.ElementTree(root)
tree.write("medusa-freq.xml", encoding="ISO-8859-1", xml_declaration = True)