#!/usr/bin/env python3
#
# Development Order #5:
#
# This is the meat and bones of the tool, where the actual desired
# commands or operation will be run. The results are then recorded
# and added to the 'results' JSON data, which will then be sent
# back to the test. Both system and api are able to be used here.
#

import os
import sys
import time
import json
import datetime
import pscheduler
import subprocess

# from stdin
input = pscheduler.json_load(exit_on_error=True)
log = pscheduler.Log(prefix='tool-curl-d2d', quiet=True)


# Take input from test spec
try:
    test_type = input['test']['type']
except KeyError as ex:
    pscheduler.fail('Missing data in input')

assert test_type == "disk-to-disk", "Internal problem."

try:
    schema = input['test']['spec']['schema']
    dest = input['test']['spec']['dest']
    source = input['test']['spec']['source']
    cleanup = input['test']['spec'].get('cleanup', True)
except KeyError as ex:
    pscheduler.fail('Missing data in input')




timeout_iso = input['test']['spec'].get('timeout', 'PT10S')

timeout = pscheduler.timedelta_as_seconds( \
               pscheduler.iso8601_as_timedelta(timeout_iso))


succeeded = False
error = ''
diags = []

STDERR = ""

# Run the actual task here:
if True:
  # Normalize input with Globus tool
  if (len(source) > 6 and source[:7] != "file://"):
    # If source is local file, use file:// protocol
    if (source[0] == "/" and source[1] != "/"):
      source = "file://" + source

  if dest and (len(dest) > 6 and  dest[:7] == "file://"):
    dest = dest.split("file://")[-1]

  # If the destination already exists, it can only be a plain file.
  # Anything else leaves open avenues for monkey business.
  if os.path.exists(dest) and not os.path.isfile(dest):
      pscheduler.succeed_json({
          'succeeded': False,
          'error': 'Destination exists but is not a plain file.'
      })

  # Execute curl command
  argv = ['curl']
  argv.append('-o') # Flag to output file
  argv.append(dest)
  argv.append(source) # Source URL
  
  try:
      start_at = input['schedule']['start']
      log.debug("Sleeping until %s", start_at)
      pscheduler.sleep_until(start_at)
      log.debug("Starting")
  except KeyError:
      pscheduler.fail("Unable to find start time in input")

  diags.append(' '.join(argv))
  start_time = datetime.datetime.now()
  status, stdout, stderr = pscheduler.run_program(argv, timeout=timeout)
  end_time = datetime.datetime.now()
  diags.append(stdout)
  diags.append(stderr)

  lines = stderr.split('\n')
  parsed = [stat for stat in lines[-1].split(' ') if stat != '']
  
  STDERR = parsed

  if status:
    succeeded = False
    error = "Error running program:\n%s"% stderr.strip('\n')

    # Determine how much (if any) of the file was transferred
    if dest and os.path.isfile(dest):
      bytes_sent = os.stat(dest).st_size

      # Convert bytes to human readable format 
      # TODO: The pscheduler module has functions for doing this.
      scalar = ['B', 'KB', 'MB', 'GB', 'TB']
      while bytes_sent > 1024 and len(scalar) > 1:
        bytes_sent = int(bytes_sent / 1024)
        scalar = scalar[1:]

      # Put this into diags
      #error += '\n\nBytes Transfered: %s %s'% (int(bytes_sent), scalar[0])
      #error += '\n\nActual Timeout: %s seconds' % timeout
      
      # Getting source filesize
      #argv = ['curl', '-s', source, '| wc -c']
      #outs = pscheduler.run_program(argv, timeout=timeout)
      #error += '\n\n' + str(outs)

  else:
    succeeded = True
    

if cleanup:
    if os.path.isfile(dest):
        try:
            os.remove(dest)
        except Exception as ex:
            diags.append("Exception while removing file: %s" % (str(ex)))


# Organize results into json data


results = {
    "succeeded": succeeded
}


results['time'] = pscheduler.timedelta_as_iso8601(end_time - start_time)


STDERR = stderr.split("\n")[-2].split()

#STDERR[6] is Average Download Speed
#Try to get download speed, if can't parse then skip
try:
    # This is in bytes/sec; convert to bits/sec.
    results['throughput'] = pscheduler.si_as_number(STDERR[6]) * 8
except Exception as ex:
    # TODO: This should probably fail since throughput is what the
    # measurement is all about.
    diags.append("Throughput speed exception: %s" % (str(ex)))


# TODO: Refactor this in the same style as the above try/except.
try:
    results['bytes-sent'] = pscheduler.si_as_number(STDERR[1])
except Exception as ex:
    diags.append("Failed to get bytes sent: %s" % (str(ex)))
    pass



pscheduler.succeed_json({
    'succeeded': succeeded,
    'result': results,
    'diags': "\n".join(diags)
})
