#!/usr/bin/python3
import subprocess
import os
import logging
import re
import sys
import uuid
from time import sleep

PACKAGE_VERSION = "A.B.X"
logger = logging.getLogger("c2v-convert")


def check_vzlist(ct_uuid):
    try:
        output = subprocess.check_output(['vzlist', '-a'], universal_newlines=True)
        logger.debug(output)
        if ct_uuid in output:
            return True
        else:
            logger.error(f"CT {ct_uuid} not found in vzlist.")
            return False
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to run vzlist: {e}")
        return False

def check_config_link(ct_uuid):
    config_path = f"/etc/vz/conf/{ct_uuid}.conf"
    if os.path.islink(config_path):
        logger.info("%s is link to the container config and therefore configs are in sync" % config_path)
        return True
    else:
        logger.error(f"Config file {config_path} is not a symbolic link.")
        return False

def check_centos_and_pp_release(ct_uuid):
    try:
        os_release = subprocess.check_output(['vzctl', 'exec2', ct_uuid, 'cat', '/etc/os-release'], universal_newlines=True)
        logger.debug(os_release)
        if 'CentOS' not in os_release:
            logger.error(f"CT {ct_uuid} is not using CentOS.")
            return False
        
        pp_release = subprocess.check_output(['vzctl', 'exec2', ct_uuid, 'rpm', '-q', 'pp-release'], universal_newlines=True)
        logger.debug(pp_release)
        if 'not installed' in pp_release:
            logger.error(f"pp-release package is not installed in CT {ct_uuid}.")
            return False
        
        logger.info("OS Version and PP check is completed successfully")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to check CentOS or pp-release: {e}")
        return False

def create_backup(ct_uuid, backup_ct_uuid):
    try:
        subprocess.check_output(['vzmlocal', '--online', '-C', f'{ct_uuid}:{backup_ct_uuid}'], universal_newlines=True)
        output = subprocess.check_output(['prlctl', 'list','-i', ct_uuid], universal_newlines=True)
        logger.debug(output)
        id=re.search(r'.*Name: (.*)\nDes.*', output)
        if id:
            subprocess.check_output(['prlctl', 'set', backup_ct_uuid, '--name', id.group(1)+"_pp_centos7_bkp"], universal_newlines=True)
        logger.info(f"Backup created successfully with BACKUP_CTUUID: {backup_ct_uuid}")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to create backup: {e}")
        return False
    return True

def convert_ct(ct_uuid):
    try:
        subprocess.check_output(['vzctl', 'set', ct_uuid, '--ostemplate', '.vzlinux-7-x86_64', '--distribution', 'vzlinux', '--save'], universal_newlines=True)
        logger.info(f"CT {ct_uuid} switched successfully to vzlinux distribution.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to convert CT {ct_uuid}: {e}")
        return False
    return True

def revert_ct(ct_uuid):
    try:
        subprocess.check_output(['vzctl', 'set', ct_uuid, '--ostemplate', '.centos-7-x86_64', '--distribution', 'centos', '--save'], universal_newlines=True)
        logger.info(f"CT {ct_uuid} reverted to centos distribution.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to revert CT {ct_uuid}: {e}")
        return False
    return True

def revert_backup(ct_uuid, backup_ct_uuid):
    try:
        output = subprocess.check_output(['prlctl', 'list','-i', ct_uuid], universal_newlines=True)
        logger.debug(output)
        id=re.search(r'.*Name: (.*)\nDes.*', output)
        if id:
            name=id.group(1)
        else:
            name=ct_uuid
        subprocess.check_output(['prlctl', 'stop', ct_uuid], universal_newlines=True)
        subprocess.check_output(['prlctl', 'set', ct_uuid, "--name",name+"-convert-failed"], universal_newlines=True)
        subprocess.check_output(['prlctl', 'set', backup_ct_uuid, "--name",name], universal_newlines=True)
        subprocess.check_output(['prlctl', 'start', backup_ct_uuid], universal_newlines=True)
        logger.info(f"CT {name} reverted to initial state from backup and started,  resulting container is left under the '{name}-convert-failed' name for investigation")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to revert CT state {name}: {e}")
        return False
    return True

def upgrade_ct(ct_uuid):
    try:
        subprocess.check_output(['vzpkg', 'update', ct_uuid], universal_newlines=True)
        logger.info(f"CT {ct_uuid} upgraded successfully.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to upgrade CT {ct_uuid}: {e}")
        return False
    return True

def run_distro_sync(ct_uuid):
    try:
        subprocess.check_output(['vzctl', 'exec3', ct_uuid, 'yum', '-y', 'distro-sync'], universal_newlines=True)
        logger.info(f"yum -y distro-sync completed successfully in CT {ct_uuid}.")
    except subprocess.CalledProcessError as e:
        logger.warning(f"Failed to run yum -y distro-sync in CT {ct_uuid}: {e}")
        return False
    return True

def run_internal_commands(ct_uuid):    
    yum_shell_command = "echo -e 'install python-oauthlib \n remove python2-oauthlib \n run \n exit' | yum shell -yy"
    try:
        subprocess.check_output(['vzctl', 'exec3', ct_uuid, 'bash', '-c', yum_shell_command], universal_newlines=True)
        logger.info(f"Yum shell commands executed successfully in CT {ct_uuid}.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to run yum shell commands in CT {ct_uuid}: {e}")
        return False
    return True

def restart_ct(ct_uuid):
    try:
        subprocess.check_output(['prlctl', 'restart', ct_uuid], universal_newlines=True)
        logger.info(f"CT {ct_uuid} have been restarted.")
    except subprocess.CalledProcessError as e:
        logger.error(f"CT {ct_uuid} failed to restart: {e}")
        return False
    return True

def check_pp(ct_uuid):
    attempt=0
    while True:
        try:
            subprocess.check_output(['prlctl', 'exec', ct_uuid, 'wget', '127.0.0.1:6556', '-q'], universal_newlines=True)
            logger.info(f"PowerPanel vzapi in CT {ct_uuid} is accessible")
            break
        except subprocess.CalledProcessError as e:
            attempt+=1
            if attempt>2:
                logger.error(f"PowerPanel vzapi in CT {ct_uuid} is inaccessible")
                return False
            sleep(10)
    return True

def check_needs_restarting(ct_uuid):
    try:
        # Run the needs-restarting command and capture the output and the exit code
        result = subprocess.run(['vzctl', 'exec3', ct_uuid, 'needs-restarting', '-r'], 
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        output = result.stdout.strip()
        exit_code = result.returncode
        
        # Print the output of the command
        logger.debug(f"Output of needs-restarting -r for CT {ct_uuid}:\n{output}")
        
        # If the output indicates that a reboot is required, suggest restarting
        if "Reboot is required" in output:
            logger.info(f"Restart is recommended for CT {ct_uuid}.")
            return True
        elif exit_code != 0:
            # Handle cases where the command fails but we still need to evaluate the output
            logger.warning(f"Command returned a non-zero exit status {exit_code}, but evaluating output.")
            if "Core libraries or services have been updated" in output:
                logger.warning(f"Restart is recommended for CT {ct_uuid} due to core updates.")
                return True
        else:
            logger.warning(f"No restart is necessary for CT {ct_uuid}.")
            return False
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to check if restart is needed for CT {ct_uuid}: {e}")
        return False

def validate_name(ct):
    try:
        status = subprocess.check_output(['vzctl', 'status', ct], universal_newlines=True)
        logger.debug(status)
        id=re.search(r'VEID (.*) exist', status)
        if id:
            return id.group(1)
        else:
            return ct
    except:
        return None



def main(ct_uuid):
    if not check_vzlist(ct_uuid):
        sys.exit(1)
    
    if not check_config_link(ct_uuid):
        sys.exit(1)

    if not check_centos_and_pp_release(ct_uuid):
        sys.exit(1)
    
    backup_ct_uuid = str(uuid.uuid4())

    if not create_backup(ct_uuid, backup_ct_uuid):
        sys.exit(1)

    if not convert_ct(ct_uuid):
        sys.exit(1)

    if not upgrade_ct(ct_uuid):
        if revert_ct(ct_uuid):
            logger.warning(f"Reverted CT {ct_uuid} to initial state.")
            sys.exit(1)
        else:
            logger.error(f"Failed to revert CT {ct_uuid} to initial state.")
            sys.exit(1)

    run_distro_sync(ct_uuid)

    if not run_internal_commands(ct_uuid):
        revert_backup(ct_uuid, backup_ct_uuid)
        sys.exit(1)

    restart_ct(ct_uuid)
    if check_pp(ct_uuid):
        logger.info(f"CT {ct_uuid} PowerPanel seems to be working")
        if check_needs_restarting(ct_uuid):
            restart_ct(ct_uuid)
    else:
        logger.error(f"CT {ct_uuid} PowerPanel is not working")
        revert_backup(ct_uuid, backup_ct_uuid)
        if check_needs_restarting(backup_ct_uuid):
            restart_ct(backup_ct_uuid)
        sys.exit(1)


    if check_needs_restarting(ct_uuid):
        restart_ct(ct_uuid)

    logger.info(f"All operations completed successfully for CT {ct_uuid} with BACKUP_CTUUID: {backup_ct_uuid}.")

if __name__ == "__main__":
    logger.setLevel(logging.DEBUG)
    log_format="%(asctime)s %(levelname)s: %(message)s"
    log_formater=logging.Formatter(log_format)
    consoleHandler = logging.StreamHandler()
    consoleHandler.setLevel(logging.DEBUG)
    consoleHandler.setFormatter(log_formater)
    logger.addHandler(consoleHandler)
    

    if len(sys.argv) != 2:
        logger.error("Wrong usage in script")
        print("Usage: ppconvert <CT_UUID | CT_NAME>")
        sys.exit(1)

    
    ct_uuid = sys.argv[1]
    logger.info("Converting %s" % ct_uuid)

    ct_uuid=validate_name(ct_uuid)
    if ct_uuid:
        main(ct_uuid)

