#!/usr/bin/python
#
# fitspipe-get,py - an attempt to validate the FitspipeProtocol
# wiki page by reading it and producing a simple working a client
# that can fetch frames.
#
# Not really intended for serious use, and probably not a great
# advert for python best proactices, either - just a quick hack job
# to validate the documentation.

import argparse # Argument parsing
import shlex    # Simplest way found for parsing name=value pairs after 2 secs googling.
import socket   # Networking
import sys      # To get at sys.args, sys.stderr, etc.
import string   # Using splitlines() to iterate through lines in response.

from enum import Enum # Will enumerate status indicators to try to make things more readable.

# Indexes the different prefixes we expect in the first character of
# a response from the server.
class MSG_PREFIX(Enum):
  CONTINUE = 1
  ERROR = 2
  GET = 3
  OK = 4

# Strings for prefixes.
PREFIX_STR = {}
PREFIX_STR[MSG_PREFIX.CONTINUE.value] = '+ '
PREFIX_STR[MSG_PREFIX.ERROR.value] = '! '
PREFIX_STR[MSG_PREFIX.GET.value] = '# '
PREFIX_STR[MSG_PREFIX.OK.value] = '. '

# Some FITS header lengths.
FITS_BLOCK_LEN   = 2880 # The length of one block of FITS header data.
FITS_CARD_LEN    = 80   # Length of a single card in the header.
FITS_KEYWORD_LEN = 8    # Keywords are this long.

###################################################################################

##########
# Sanity-checking of command line args, used by parse_args(). Returns
# boolean indication of valid args or not.
def check_args(args):
  if(args.port < 2) or (args.port > 65535):
    sys.stderr.write("error: Port %d is invalid (%d..%d)\n" % (args.port, 1, 65535))
    return False

  # Not interested in supporting ridiculously large captures - this
  # is just a proof of concept.
  #
  # We could support infinitely long streams by perhaps allowing nframe to
  # be None, and then looping accordingly in get_frames().
  if(args.nframe < 1) or (args.nframe > 1000):
    sys.stderr.write("error: Number of frames %d is invalid (%d..%d)\n" % (args.nframe, 1, 100))
    return False

  return True


##########
# Python's argparse module uses args that are of the form
# --argname value rather than our name=value pairs that we
# tend to like. Still, it's not the point of this to worry
# about that, so using what's quickly and readily available
# to create something readable.
def parse_args(argv):
  parser = argparse.ArgumentParser()

  parser.add_argument('server',
                      help='Hostname/address of fitspipe server')

  parser.add_argument('--port',
                      dest='port',
                      required=False,
                      type=int,
                      default=9999,
                      help='Port number for fitspipe server [9999]')

  parser.add_argument('--feed',
                      dest='feed',
                      required=False,
                      default='default',
                      help='fitspipe feed to fetch frame(s) from [default]')

  parser.add_argument('--nframe',
                      dest='nframe',
                      required=False,
                      type=int,
                      default=1,
                      help='Number of frames to retrieve [1]')

  parser.add_argument('--out',
                      dest='out',
                      required=False,
                      default=None,
                      help='Partial path to write FITS files to. Leave out to write to stdout')

  args = parser.parse_args()

  if check_args(args) != True:
    return None
  else:
    return args


##########
# Takes a string containing whitespace-separated name=value
# parameters, turns it into a dict.
def dict_from_name_value_pairs(s):
  lexer = shlex.shlex(s, posix=True)
  lexer.whitespace_split = True
  lexer.whitespace = ' '
  return dict(pair.split('=', 1) for pair in lexer)


##########
# Turns one line of response from 'ls' command into something
# usable. Returns a dict containing all the info about the
# feed, with key names matching those from the 'ls' command
# (e.g. feed name is 'feed', etc.).
#
# Does not validate the magic status character at the head
# of the line. Assumes that the response has already been
# checked for sanity.
def parse_feed(line):
  # Parse everything after the message prefix.
  return dict_from_name_value_pairs(line[2:])


##########
# Parses the prefix on the head of a line of response from
# the server. Returns an entry from the MSG_PREFIX enumeration
# if that status indication was found, otherwise None if there
# was nothing sensible.
def get_resp_prefix(line):

  for prefix in MSG_PREFIX:
    if line[:2] == PREFIX_STR[prefix.value]:
      return prefix

  # Nothing found.
  return None


##########
# Gets info for the requested feed from the command line.
# Returns a dict representing feed, with keys with names
# and values that match what came out of the server.
def get_feed(args, sock):

  # Server will respond to 'ls' commandsends info on all feeds, one per line.
  msg='ls'
  sock.send(msg.encode())

  # In the real world, we'd possibly worry about receiving
  # large lists of feeds and would do something more sensible
  # involving receiving multiple lines, etc. For this little
  # charade, we'll just assume# an upper limit that's far larger
  # than anything we're likely to see in a simple system.
  # Note: response decoded as ascii text so we can treat it as a
  # string.
  resp = sock.recv(10000).decode("utf-8")

  # Try to match the feed we were given on the command line
  # with whatever the server reckons it has.
  found_feed = None
  for line in resp.splitlines():

    # Get prefix from message.
    prefix = get_resp_prefix(line)

    # Still looking for named feed.
    if found_feed is None:

      # Explicit error?
      if prefix == MSG_PREFIX.ERROR:
        sys.stderr.write("error: Server returned error\n")
        return None

      # If the server indicates success, then we didn't
      # find a matching feed.
      elif prefix == MSG_PREFIX.OK:
        return found_feed

      # If we hit anything but one of the continuation lines,
      # then that's also an error,
      elif prefix != MSG_PREFIX.CONTINUE:
        sys.stderr.write("error: Server returned unexpected line prefix\n")
        return {'feed':'!notfound'}

      # Figure out if this is the feed we're looking for.
      feed = parse_feed(line)
      if feed['feed'] == args.feed:
        found_feed = feed

    else:
      # Found our feed. Waiting for an indication of success.
      if prefix == MSG_PREFIX.OK:
        if found_feed is None:
          sys.stderr.write("No feed '%s' found on server.\n" % args.feed)
        return found_feed

  # If we get here, then something went wrong in the response from the server.
  sys.stderr.write("error: Unexpected response to 'ls' command from server.\n")
  return None


##########
# Takes a 40-byte initial response to 'get' and makes sense of it.
# Strongly based on the assumption that the response has a very
# rigid format - if this ever changes on the server side, this will
# likely break horribly.
def parse_get_header(line):
  info = {'frame':0, 'width':0, 'height':0, 'nbytes':0}

  # Anything but the magic prefix on the head of the response is an error.
  prefix = get_resp_prefix(line)
  if prefix != MSG_PREFIX.GET:
    sys.stderr.write("error: Unexpected response from server '%s'\n" % line)
    return None

  # Fields are guaranteed to be in fixed positions in the string.
  info['frame'] = int(line[2:13])
  info['width'] = int(line[13:23])
  info['height'] = int(line[26:36])

  # 16bpp pixels.
  info['nbytes'] = info['width'] * info['height'] * 2

  return info

##########
# Receives a fixed amount of data (n bytes) from the socket (s),
# blocking until all data received (hopefully).
#
# Used only when the amount of data expected back from the server
# is known ahead of time (i.e. the 40-line intiial response to 'get',
# blocks of FITS headers, or the image data).
#
# Since python doesn't provide any standard way of completely receiving
# a fixed amount of data and I couldn't be bothered writing my own,
# this was stolen from:
#
#   https://stackoverflow.com/questions/17667903/python-socket-receive-large-amount-of-data
#
# Not exactly guaranteed to be completely safe in the face of various
# problems, but that's not the point of this example, which demonstrates
# receiving fitspipe frames. If you want robust networking, rip out the
# socket handling side of this tool and replace it with something more
# completely thought-out.
def recvall(sock, n):
    # Helper function to recv n bytes or return None if EOF is hit
    data = b''
    while len(data) < n:
        packet = sock.recv(n - len(data))
        if not packet:
            return None
        data += packet
    return data


##########
# Given a block of FITS header data, this indicates whether this is
# the last block of the header or not. Returns boolean to show whether
# this is the end of the header.
def last_fits_header_block(block):
  i = 0

  # Walk through header, 80 bytes at a time, looking for an 'END' keyword
  # that signifies the end of the header. Only look at the beginning of
  # 80-byte boundaries so we don't get spoofed by other arbitrary strings,
  # and look for trailing spaces in keyword so we don't get caught out by
  # other keywords that start with 'END'.
  while i < FITS_BLOCK_LEN:
    if block[i:i+FITS_KEYWORD_LEN] == "END     ":
      return True

    i += FITS_CARD_LEN

  return False


##########
# Gets a full FITS header from the server without actually parsing it.
# The trick here is that we don't know ahead of time how large the header
# is going to be, but with a few simple FITS rules, we can figure it out
# as we go along.
def get_fits_header(sock):
  header= b''

  # Repeatedly get blocks of FITS header until we find the last one.
  # We'd probably do something more intelligent in a 'real' fitspipe
  # tool, but the length check here is just to bail if we've done
  # something inadvisable and not get stuck in an infinite loop.
  while len(header) < (100 * FITS_BLOCK_LEN):

    # Get next block.
    block = recvall(sock, FITS_BLOCK_LEN)
    if block is None:
      sys.stderr.write("error: Failed to get FITS header block from server.\n")
      return None

    # Add this block to the header received so far.
    header += block

    # Is this the last block?
    if last_fits_header_block(block):
      return header

  # If we get here, either the header is weirdly large or we did
  # something very wrong.
  sys.stderr.write("error: Getting FITS header took unexpectedly long. Bailing.\n")
  return None


##########
# Gets one or more frames from server. 'feed' is the
# info retreived from the server for the feed with 'ls' ahead
# of the transfer.
def get_frames(args, feed, sock):

  i = 1
  while i <= args.nframe:
    next_frame = int(feed['newest']) + i #- int(feed['depth'])

    # Request the next frame from the server.
    msg = 'get feed=' + feed['feed'] + ' frame=' + str(next_frame) + ' fullheader=1'
    sock.send(msg.encode())

    # Get the first line response from the server. This is a fixed-length
    # message.
    resp = recvall(sock, 40).decode("utf-8")
    image_info = parse_get_header(resp)

    if image_info is None:
      return

    # Get the FITS header from the server ahead of the image data.
    header = get_fits_header(sock)
    if header is None:
      return

    # Get expected image size from server.
    sys.stderr.write("Frame %d: getting %d bytes\n" %(next_frame, image_info['nbytes']))
    image = recvall(sock, image_info['nbytes'])

    if image is None:
      sys.stderr.write("error: Failed to receive data for frame %d\n" % (next_frame))
      return

    # FITS standard requires that we pad out to the end of the
    # nearest block size with zeroes. Figure out how many bytes
    # are going to be needed.
    if (image_info['nbytes'] % FITS_BLOCK_LEN) != 0:
      padding_len = FITS_BLOCK_LEN - (image_info['nbytes'] % FITS_BLOCK_LEN)

      # Generate padding, helpfully zeroed out and ready to go.
      padding = bytearray(padding_len)
    else:
      padding = None

    # FITS frame either gets chucked on stdout, or it ends up in a file.
    if args.out is None:
      sys.stdout.write(header)
      sys.stdout.write(image)
      if padding is not None:
        sys.stdout.write(padding)

    else:
      # Padding frame number with leading zeroes in filename.
      filename = args.out + str(i).zfill(5) + ".fits"
      f = open(filename, "w")
      f.write(header)
      f.write(image)
      if padding is not None:
        f.write(padding)
      f.close()

    # Next frame.
    i+=1


###################################################################################

# Parse command line.
args = parse_args(sys.argv)
if args is None:
  exit(1)

# Create a socket and connect to the server. Note complete lack
# of error checking here - that's not the point of this example.
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect((args.server, args.port))

# Get info for the requested feed.
feed = get_feed(args, s)

# If we found a feed, then we're OK to go get frame(s) from
# the server.
if feed is not None:
  get_frames(args, feed, s)

# Done.
s.close()