#!/usr/bin/python3
import json
import re
import subprocess
from subprocess import Popen, PIPE
import time
import datetime
import os
from os import path
import sys
from collections import defaultdict
from collections import OrderedDict

try:
    from viavi.mts.device_th import send_od_to_otu_logging
except ImportError:
    from jdsu.mts.device_th import send_od_to_otu_logging
    
##############################################
def extract_upgrade_version(fp_name):
     otu_upgrade_file = fp_name
#     version_pattern = "otu\-5000\-(.*?)\.tar"
     if ("otu-5000" in fp_name):
         version_pattern = "otu\-5000\-(.*?)\.tar"
     elif ("fth-9000" in fp_name):
         version_pattern = "fth\-9000\-(.*?)\.tar"
     elif ("fth-7000" in fp_name):
         version_pattern = "fth\-7000\-(.*?)\.tar"
     else:
         version_pattern = "otu\-8kv2\-(.*?)\.tar\.gz"
     upgrade_version_number = re.search(version_pattern, otu_upgrade_file).group(1)
     return (upgrade_version_number)
     
############## >>>>>>>>>>
#  Verify the path of the upgrade tarball file
def verify_upgrade_file_exists(upg_file_path):
    file_path = path.exists(upg_file_path)
    if file_path != True:
        #print ("Input file ", upg_file_path, "does not exist")
        rc = False
    else:
        #print ("Input file is available at ", sys.argv[1])
        rc = True
    return rc

############ >>>>>>>>>>>>
def print_to_console(status, previous_version, upgrade_version, duration, upgrade_time):
    ztp_upg_time = "OTU ZTP Upgrade Info ----  " + upgrade_time 
    mnf1 = "{"
    mnf2 = "    \"upgrade_status\": " + "\"" + status + "\","
    mnf3 = "    \"previous_version\": " + "\"" + previous_version + "\","
    mnf4 = "    \"current_version\": " + "\"" + upgrade_version + "\","
    mnf5 = "    \"duration\": " + "\"" + str(duration) + "\""
    mnf6 = "}"
    lnspc = "\n"
    #sys.stdout.write(ztp_upg_time+"\n"+mnf1+"\n"+mnf2+"\n"+mnf3+"\n"+mnf4+"\n"+mnf5+"\n"+ mnf6+lnspc)
    sys.stdout.write(mnf1+"\n"+mnf2+"\n"+mnf3+"\n"+mnf4+"\n"+mnf5+"\n"+ mnf6+lnspc)
    sys.stdout.flush()
    return (True)

############ >>>>>>>>>>>>
def log_upgdata(con_error, status, previous_version, upgrade_version, duration, upd_time):
    upg_info = OrderedDict((("upgrade_status" , status), \
                 ("previous_version", previous_version), \
                 ("upgrade_version", upgrade_version), \
                 ("duration", str(duration))))
    ztp_conf_upd_time = "OTU ZTP Upgrade Info ----  " + upd_time
    send_od_to_otu_logging(con_error, ztp_conf_upd_time, upg_info)
    return True

if __name__ == "__main__":
    start_time = time.time()
    err_details = ""
    e_code=0
    failed_status = "failed"
    success_status = "success"
    upgrade_status = "initiated OTU upgrade"
    previous_version = "not available"
    upgrade_version = "not available"

    reboot_delay = 15
    max_python_script_inputs = 2
    max_input_arguments = max_python_script_inputs - 1
    otu_ztp_upgrade_time = time.strftime("%c")

    #
    #  Get the current (previous) running version of the OTU
    otu_info_text = (((subprocess.check_output(["/sbin/getinfoversion", "Otu"])).decode("utf-8")).split(";"))
    previous_version = otu_info_text[3].replace("V",'')
    #
    #  Validate the the number of input arguments
    groot_arguments = len(sys.argv)
    if groot_arguments != max_python_script_inputs:
        #
        # Create error detail message err_detail
        err_details = "ZTP Upgrade Input Error --- number of input arguments must be " + str(max_input_arguments) 
        upgrade_status = failed_status
        #sys.stdout.write(err_details + "\n")
        #sys.stdout.flush()
        #
        # Calculate upgrade duration, set the update status and log the data
        upgrade_dur = (datetime.timedelta(seconds=(time.time() - start_time)))
        print_to_console(upgrade_status, previous_version, upgrade_version, upgrade_dur, otu_ztp_upgrade_time)
        log_upgdata(err_details, upgrade_status, previous_version, upgrade_version, upgrade_dur, otu_ztp_upgrade_time)
        #  Set the error code and exit the script
        e_code = 1
        exit(e_code)
    #
    #  Validate the path and presence of the ZTP upgrade tarball file
    groot_path = verify_upgrade_file_exists(sys.argv[1])
    if groot_path != True:
        # Calculate upgrade duration, set status and log the meaningful data
        upgrade_dur = (datetime.timedelta(seconds=(time.time() - start_time)))
        upgrade_status = failed_status
        err_details = "ZTP Upgrade Error --- input file " + sys.argv[1] + " does not exist"
        #sys.stdout.write(err_details + "\n")
        #sys.stdout.flush()
        print_to_console(upgrade_status, previous_version, upgrade_version, upgrade_dur, otu_ztp_upgrade_time)
        log_upgdata(err_details, upgrade_status, previous_version, upgrade_version, upgrade_dur, otu_ztp_upgrade_time)
        #  Set error code and exit the script
        e_code = 1
        exit (e_code)
    #
    #  Fretch the new upgrade version number of the OTU upgrade
    upgrade_version = extract_upgrade_version(sys.argv[1])
    #
    #  Start the OTU upgrade
    upg_request = "/sbin/start_upgrade_otu.sh --tar " + sys.argv[1] + " --nodet --noreboot --target NORMAL --index 0 --nomd5"
    #sys.stdout.write("Upgrade is in progress for " + sys.argv[1] + "\n")
    groot_cmd = subprocess.Popen(upg_request, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, universal_newlines=True)
    std_out, std_err = groot_cmd.communicate()
    groot_rc = groot_cmd.returncode
    #  Calculate ZTP upgrade duration
    upgrade_dur = (datetime.timedelta(seconds=(time.time() - start_time)))
    #  set upgrade status indicator
    if (groot_rc == 0):
        #err_details = "OTU ZTP upgrade completed successfully"
        upgrade_status = success_status
        e_code = 0
    else:
        #err_details = "OTU ZTP upgrade failed"
        upgrade_status = failed_status
        e_code = 1
    #  Write ZTP upgrade info to console and also log it in /var/log/user.log
    print_to_console(upgrade_status, previous_version, upgrade_version, upgrade_dur, otu_ztp_upgrade_time)
    log_upgdata("", upgrade_status, previous_version, upgrade_version, upgrade_dur, otu_ztp_upgrade_time)
    #
    #  Reboot the OTU if the upgrade was successfull)
    #  Allow 15 second wait period before rebooting
    if ( groot_rc == 0):
        #sys.stdout.write("OTU will reboot in %s seconds" % reboot_delay)
        #sys.stdout.flush()
        time.sleep(reboot_delay)
        subprocess.call("/sbin/reboot")
    #
    exit (e_code)
