#!/usr/bin/env python3

import os
import argparse
import json
import xmlrpc.client
import ssl
import time
import sys

from pathlib import Path
from n4d import client
'''
USAGE:

ltsp-lliurex-update [-i image ] [ --mode { update | regenerate} ] [--update-mirror]

'''


def update_image(image):
    try:
        print("\nltsp-lliurex-update: Updating chroot... "+image+"\n")
        cmd="llx-chroot /opt/ltsp/{image} lliurex-upgrade -u".format(image=image)
        os.system(cmd)
        return 0
        
    except Exception as e:
        print("\n\n Exception Upgrading chroot !! \n\n ")
        print(str(e))
        return 1
        
def regenerate_image(image):
    try:
        print("\nltsp-lliurex-update: Checking chroot is completely umounted for {image}\n".format(image=image))
        print("\nltsp-lliurex-update: Updating Kernels for image {image}\n".format(image=image))
        cmd = "ltsp kernel {image}".format(image=image)
        #print "Regenerate ... "+image
        os.system(cmd)
        
        cmd="ltsp initrd"
        os.system(cmd)
        
        return 0
    
    except Exception as e:
        print("\n\n EXCEPTION UPDATING IMAGE !! \n\n ")
        print(str(e))
        return 1
        
    
    pass


def update_mirror():
    c = client.Client("https://localhost:9779")
    try:
        response = c.MirrorManager.get_available_mirrors()
    except client.CALL_FAILED as e:
        print(e.message)
        sys.exit(1)
    except Exception as e:
        print(str(e))
        sys.exit(1)

    for mirror in response:
        print("mirror")
        print("Updating mirror for "+mirror+".")
        c.MirrorManager.update("","",mirror)
        while True:
            time.sleep(3)
            print('.')
            try:
                response = c.MirrorManager.is_alive()
                if (not response['status']):
                    break
            except Exception as e:
                print("!!Failed to update mirror {mirror}".format(mirror=mirror))
                print(str(e))
                sys.exit(1)
    return 0

def ltsp_update(args):
    
    confdir = Path( "/etc/ltsp/images/" )
    globalErr = 0
    
    if ( args.update_mirror ):
        globalErr = update_mirror()
            
    
    for i in confdir.iterdir():
        if i.suffix == ".json":
            try:
                with i.open("r", encoding="utf-8") as fd:
                    data = json.load(fd)

                localError=0
                if ( args.i == None or args.i == data["id"] ):
                    if (args.mode==None or args.mode=="update"):
                        localError=update_image(data["id"])
                        
                    if (localError==0 and (args.mode==None or args.mode=="regenerate")):
                        # If there where errors in update, not regenerate img
                        regenerate_image(data["id"])

            except Exception as e:
                print("Exception "+str(e))
        
    if (globalErr==0):
        print("\n\nltsp-lliurex-update: OK. Updates finished without errors")
    else:
        print("\n\nltsp-lliurex-update: FAIL.Updates finished with errors!!")

    pass


def main():
    parser = argparse.ArgumentParser(description='Updates one or more images in a ltsp server.')
    parser.add_argument('-i', help='Specifies an image. If none, all are selected.')
    parser.add_argument('--mode', help='Specifies how image/s are updated (update, regenerate). If none specified, performs all operations.', choices=["update", "regenerate"])
    parser.add_argument('--update-mirror', help='Updates mirror before update clients', action='store_true', default=False)

    args = parser.parse_args()

    ltsp_update(args)



if __name__ == "__main__":
    main()