#!/usr/bin/env python
# coding=utf-8
"""
Script to relocate the space-time bounding box of one or more existing (‘template’) IDV bundles
and then execute runIDV to fetch the data and write ouputs.

more info: https://github.com/suvarchal/IDV_teleport

"""
import argparse
import os
import re
import sys
import subprocess
from distutils import spawn
from random import randint
from zipfile import ZipFile
from tempfile import NamedTemporaryFile
import base64
from xml.etree import ElementTree
import requests


def parse_date_time_file(datefile):
    """
    :param: takes in a file object containing dates and times
    :return: sting of dates and times

    TODO: check if the times are valid or not and return only valid times.
    """
    datelist = []
    with open(datefile) as fil:
        for line in fil.readlines():
            checkline = line.strip().split('-')
            assert len(checkline) > 1, "Time specified in input is not in valid format yyyy-mm-dd"
            datelist.append(line)
    return datelist


def isl_string(bundle_file, time_start, time_end, case_name,
               ul_lat=None, ul_lon=None, lr_lat=None, lr_lon=None):
    """
    creates an isl string for IDV to run;
    :param bundle_file : is a template bundle
    :return : isl string
    """
    screencapture = u"""
def screencapture(width=None,height=None):
    import java
    import java.awt.Robot as Robot
    import java.awt.Rectangle as Rectangle
    import java.awt.Toolkit as Toolkit
    from ucar.unidata.util import Misc
    from java.awt import Point
    from java.awt import Toolkit
    from java.awt import GraphicsEnvironment
    from java.awt import Image
    from java.awt.image import BufferedImage
    from ucar.unidata.ui import ImageUtils

    VM=idv.getViewManager()
    VM.displayWindow.toFront()
    VMC=VM.getContents()
    VMCC=VMC.getComponent(1) # the view and legend ; 0 is left most part of view window with controls for perspective views

    gc= VMCC.getGraphicsConfiguration()

    loc = VMCC.getLocationOnScreen()
    siz = VMCC.getSize()

    # could also capture window by using window
    # DW=VM.getDisplayWindow()
    # AW=DW.getActiveWindow()
    # active window can be made fullscreen
    # gcd=gcd.setFullScreenWindow(AW.getWindow()) #gcd=graphic control device
    # quality can be improved by setting full screen or unlimitedly
    # on virtual screens by setting big screen size
    loc.x -= gc.getBounds().x
    loc.y -= gc.getBounds().y


    robotx=Robot() #gc.getDevice())


    VM.displayWindow.toFront()
    W=VM.getDisplayWindow().getActiveWindow().getWindow()
    W.setAlwaysOnTop(True)
    Misc.sleep(250)
    img=robotx.createScreenCapture(Rectangle(loc.x, loc.y,siz.width, siz.height))
    W.setAlwaysOnTop(False)
    #img=makeColorTransparent(img,VM.getBackground())
    if width != None and height != None:
        #img=ImageUtils.toBufferedImage(img.getScaledInstance(width,height,Image.SCALE_AREA_AVERAGING),
        #                                               BufferedImage.TYPE_INT_RGB)
        return ImageUtils.toBufferedImage(ImageUtils.resize(img,width,height),BufferedImage.TYPE_INT_RGB)
    else:
        return img

def screen_image(width=None,height=None):
    from ucar.unidata.ui import ImageUtils
    VM=idv.getViewManager()
    VM.toFront()
    anim=VM.getAnimation()
    anim.setCurrent(len(VM.getAnimationTimes())/2)
    img = screencapture(width,height)
    ImageUtils.writeImageToFile(img,'{0:s}'+'.png',1.0)

def screen_animation(width=None,height=None):
    from ucar.unidata.ui import AnimatedGifEncoder
    from ij import ImagePlus
    VM=idv.getViewManager()
    VM.toFront()
    e=AnimatedGifEncoder()
    e.start('{0:s}'+'.gif')
    e.setRepeat(0) # Repeat indefinely -1 for once
    anim=VM.getAnimation()
    VM.getAnimationWidget().gotoBeginning()
    for t in range(len(VM.getAnimationTimes())):
        anim.setCurrent(t)

        data=screencapture()
        e.addFrame(ImagePlus(str(t),data))
    """.format(case_name)  # ideally make them part of function call
    if not ul_lat and not ul_lon and not lr_lat and not lr_lon:
        xidv_string = u"""<isl>
                  <bundle file="{0:s}" timedriverstart="{1:s}" timedriverend="{2:s}" />
                  <pause seconds="40"/>
<displayproperties display="class:ucar.unidata.idv.control.ColorPlanViewControl">
<property name="DisplayAreaSubset" value="true"/>
</displayproperties>
                  <save file="{3:s}.zidv"/>
                  <pause seconds="60"/>
                  <pause/>
                  <jython><![CDATA[
                  {4:s}
                  ]]>
                  </jython>
                  <pause/>
                  <jython code="idv.waitUntilDisplaysAreDone()"/>
                  <jython code="screen_image(600,400)"/>
                  <jython code="screen_animation(600,400)"/>
                  <jython code="exit()"/>
                  </isl>""".format(bundle_file, time_start, time_end, case_name,
                                   screencapture)
    else:
        xidv_string = u"""<isl>
                  <bundle file="{0:s}" timedriverstart="{1:s}" timedriverend="{2:s}" bbox="{3:s},{4:s},{5:s},{6:s}"/>
                  <pause seconds="40"/>
<displayproperties display="class:ucar.unidata.idv.control.ColorPlanViewControl">
<property name="DisplayAreaSubset" value="true"/>
</displayproperties>
                  <save file="{7:s}.zidv"/>
                  <pause seconds="60"/>
                  <pause/>
                  <jython code="idv"/>
                  <jython code="idv.waitUntilDisplaysAreDone()"/>
                  <jython code="idv.viewManager.idvUIManager.closeHelpTips()"/>
                  <jython><![CDATA[
                  {8:s}
                  ]]>
                  </jython>
                  <jython code="screen_image(500,400)"/>
                  <jython code="screen_animation(500,400)"/>
                  <jython code="exit()"/>
                  </isl>""".format(bundle_file, time_start, time_end,
                                   str(ul_lat), str(ul_lon), str(lr_lat), str(lr_lon),
                                   case_name, screencapture)
    # ideally make above line not to repeat case_name
    return xidv_string


def parse_date_time(datetimelist, parser_args):
    """
    :rtype : tuple
    :param datetimelist: date and time as a string
    :param parser_args: parser args
    :return: tuple of lists with startdates,enddates,centerdates and dates
             datetime module couldn't parse.
    """
    import datetime

    match = re.match(r"(\d+)(\w+)", parser_args.timedelta)
    try:
        timedelta, time_str = match.groups()
        dtime = datetime.timedelta(**{time_str: int(timedelta)})
    except TypeError:
        print("Please set -td or --timedelta as 1seconds or 1days....")
        sys.exit(2)

    startdates = []
    enddates = []
    ignorelist = []
    centerdates = []
    for time in datetimelist:
        try:
            if len(time.strip()) > 10:
                time_c = datetime.datetime(int(time[0:4]), int(time[5:7]), int(time[8:10]),
                                           int(time[11:13]), int(time[14:16]), int(time[17:19]))
                time_s = time_c - dtime
                time_e = time_c + dtime
            else:
                time_c = datetime.datetime(int(time[0:4]), int(time[5:7]), int(time[8:10]), 0, 0, 0)
                time_s = time_c - dtime
                time_e = time_c + dtime
            startdates.append(time_s.strftime('%Y-%m-%d %H:%M:%S'))
            enddates.append(time_e.strftime('%Y-%m-%d %H:%M:%S'))
            centerdates.append(time_c.strftime('%Y-%m-%d-%H-%M-%S'))
        except (TypeError, ValueError):
            ignorelist.append(time.strip())
    return startdates, enddates, centerdates, ignorelist


def run_xvfb():
    """ Runs Xvfb if available and puts it in the background.
    :return subprocess spawned process or None if errors
    TODO: dont make Xvfb default?
    TODO: exporting display variable needs to be handled consistently with main function
    so that this is risk free function, ie.., now it sets a display variable which main changes
    back to original.
    """
    xvfb_executable = spawn.find_executable('Xvfb')
    r_int = randint(10, 99)
    xvfb_executable += ' :' + str(r_int)
    xvfb_executable += ' -screen 0 1800x1600x24+32'
    # print(xvfb_executable)
    xvfb_proc = subprocess.Popen(xvfb_executable.split())  # looks always will return
    # print(xvfb_proc.pid)
    if not xvfb_proc.poll():
        os.environ['DISPLAY'] = ':' + str(r_int) + '.0'
        return xvfb_proc
    else:
        return None


def publish(parentid, case_names):
    """ Given a RAMADDA parent_id and list of case_names this function
    publishes to ramadda with user credentials taken from environment variables
    RAMADDA_USER and RAMADDA_PASSWORD

    TODO: lot of redundant code in for loops checking if files exist
    TODO: do it with xml not by strings
    TODO: currently creating a folder is not possible, user has to supply a folder
    but can be fixed by adding an entry to xml_string group, but needs to handle
    user running script  multiple times.
    """
    assert isinstance(parentid, str), "publish_url must be a list"
    assert isinstance(case_names, list), "case_names must be a list"
    try:
        postadd = os.environ['RAMADDA']
    except KeyError:
        postadd = "https://weather.rsmas.miami.edu/repository/entry/xmlcreate"

    xml_string = ''  # '<entries>'  # add group here later
    for case_name in case_names:
        if os.path.isfile(case_name + '.zidv') and os.path.isfile(case_name + '.gif'):
            file_prefix = os.path.split(case_name)[-1]
            xml_string += '<entry name="{0}" file="{0}.zidv" ' \
                          'type="type_idv_bundle">'.format(file_prefix)
            xml_string += '<metadata inherited="false" type="content.thumbnail">'
            xml_string += '<attr fileid="{0}.gif" index="1">'.format(file_prefix)
            encoded_case_str = base64.b64encode(file_prefix + '.gif').decode('ascii')
            xml_string += '<![CDATA[{0}]]>'.format(encoded_case_str)
            xml_string += '</attr>'
            xml_string += '</metadata>'
            xml_string += '</entry>'
    # xml_string += '</entries>'

    try:
        user = os.environ['RAMADDA_USER']
        password = os.environ['RAMADDA_PASSWORD']
    except KeyError as err:
        print('Publish error {0}'.format(err))

    with NamedTemporaryFile(suffix=".zip") as tmpzip:
        with ZipFile(tmpzip.name, 'w') as zipfile:
            zipfile.writestr('entries.xml', xml_string)
            for case_name in case_names:
                if os.path.isfile(case_name + '.zidv') and os.path.isfile(case_name + '.gif'):
                    file_prefix = os.path.split(case_name)[-1]
                    zipfile.write(file_prefix + '.zidv')
                    zipfile.write(file_prefix + '.gif')
            files = {"file": open(tmpzip.name, "rb")}
            resp = requests.post(postadd, files=files,
                                 data={'group': parentid,
                                       'auth.user': user, 'auth.password': password,
                                       'response': 'xml'})
    print(tmpzip.name)
    publish_attrib = ElementTree.fromstring(resp.text).attrib
    if publish_attrib['code'] == 'ok':
        print('Published case {0} '.format(case_name))
    else:
        print('Publish case {0} failed with {1}'.format(case_name, publish_attrib['code']))
    return None


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Script to teleport time and space domain '
                                                 'of an IDV Bundle.',
                                     epilog="Simplest use case:python IDV_teleport.py "
                                            "-t timefile "
                                            "-b templatebundlefile.xidv")

    parser.add_argument('-b', '--bundle', nargs='+',
                        help='IDV Bundle template file, local or remote',
                        required=True)

    parser.add_argument('-bbox', '--boundingbox', nargs=4, type=float,
                        help='Set the bounding box of the bundle with boundaries: '
                             'north, west, south, east',
                        metavar=("NORTH", "WEST", "SOUTH", "EAST"))

    parser.add_argument('-t', '--time', type=str,
                        help='Input time as YYYY-MM-DD or a text file with times '
                             'one per line also optionally with hh:mm:ss',
                        required=True)
    parser.add_argument('-td', '--timedelta',
                        help='Time delta (hald-duration) as Nseconds, Ndays, Nweeks ...;'
                             ' output bundle times will be central time +- timedelta;'
                             ' default is 0seconds',
                        default="0seconds", required=False)

    parser.add_argument('-case', '--case_name', type=str, nargs='+',
                        help='Case name to prefix the bundle;'
                             'By default case name will be selected from bundle file',
                        required=False)
    parser.add_argument('-outdir', '--output_directory', type=os.path.isdir,
                        help='Set the output path to place the output;'
                             'default is current directory from where the script is run',
                        required=False)
    parser.add_argument('-pubid', '--publish_id',
                        help='Publish bundle and image at a RAMADDA server;'
                             'argument shoud be ramadda entryid where'
                             'the user from environment variable RAMADDA_USER'
                             'and password from RAMADDA_PASSWORD has permissions'
                             'to write files',
                        required=False)
    parser.add_argument('-nohead', '--headless', choices=("True", "False"), default="False",
                        help='Option to use headless display environment or not'
                             'to use headless environment `Xvfb` needs to be installed'
                             'and be present in the PATH.'
                             'Default is set to True for convinience,'
                             'When True and Xvfb is not on path it tries '
                             'to run IDV with default local display',
                        required=False)
    parser.add_argument('-d', '--debug', choices=("True", "False"), default="False",
                        help='Debug option; '
                             'For each time in timefile IDV session will ONLY close manually',
                        required=False)
    pargs = parser.parse_args()
    try:
        idv_home = os.environ['IDV_HOME']
    except KeyError:
        parser.print_usage()
        print("Please set environment variable IDV_HOME to IDV home directory")
        sys.exit(2)

    headless = eval(pargs.headless)
    xvfb = None

    if headless:
        orig_display_id = os.environ['DISPLAY']
        xvfb = run_xvfb()  # changes Display env variable
        if not xvfb and not orig_display_id:
            raise Exception('Headless server environment detected and Xvfb not working,'
                            ' please check if Xvfb is already running ')
    # check if datefile is time or file
    if os.path.isfile(pargs.time):
        datelist = parse_date_time_file(pargs.time)
    else:
        datelist = [str.join(' ', pargs.time.split('_'))]

    startdates, enddates, centerdates, ignorelist = parse_date_time(datelist, pargs)

    if pargs.output_directory:
        output_directory = pargs.output_directory
    else:
        output_directory = os.getcwd()

    bundle_file = r'\ '.join(pargs.bundle)
    case_names = []
    for start, end, center in zip(startdates, enddates, centerdates):
        if pargs.case_name:
            case_name = os.path.join(output_directory, pargs.case_name[0])  # +'_'+center)
        else:
            case_name = os.path.join(output_directory,
                                     os.path.split(bundle_file)[-1].split('.')[0] + '_' + center)
        if pargs.boundingbox:
            isl = isl_string(bundle_file, start, end, case_name,
                             pargs.boundingbox[0], pargs.boundingbox[1],
                             pargs.boundingbox[2], pargs.boundingbox[3])
        else:
            isl = isl_string(bundle_file, start, end, case_name)

        wtf = NamedTemporaryFile(mode="w", suffix=".isl", dir="./")
        wtf.file.write(isl)
        wtf.file.close()
        try:
            subprocess.call([os.path.join(idv_home, "runIDV"), "-islinteractive",
                             "-noerrorsingui", wtf.name.split('/')[-1]])
            case_names.append(case_name)
        except Exception:
            # not properway should check what subprocess returns
            print("Could not process the case {0}".format(case_name))
    if pargs.publish_id:
        publish(pargs.publish_id, case_names)
    if xvfb:
        if not xvfb.poll():
            xvfb.kill()
        os.environ['DISPLAY'] = orig_display_id  # imp otherwise headless
        # will not work next time
