#!/usr/bin/python3
import subprocess
import os
import json
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_name):
    try:
        output = subprocess.check_output(['vzlist', '-a','-H', ct_name], universal_newlines=True)
        logger.debug(output)
        if 'invalid' in output:
            logger.error(f"CT {ct_name} not found in vzlist.")
            return False
        else:
            return True            
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to run vzlist: {e}")
        return False
    
def get_ct_id(ct_name):
    output = subprocess.check_output(['vzlist', '-j', '-H', '-n', ct], universal_newlines=True)
    return json.loads(output)[0]["ctid"]
    

def check_config_link(ct_name):
    config_path = f"/etc/vz/conf/{ct_name}.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.warning(f"Config file {config_path} is not a symbolic link.")
        ctid=get_ct_id(ct_name)
        config_path = f"/etc/vz/conf/{ctid}.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.warning(f"Config file {config_path} is not a symbolic link.")
            return False

def check_centos_and_pp_release(ct_name):
    try:
        os_release = subprocess.check_output(['vzctl', 'exec2', ct_name, 'cat', '/etc/os-release'], universal_newlines=True)
        logger.debug(os_release)
        if 'CentOS' not in os_release:
            logger.error(f"CT {ct_name} is not using CentOS.")
            return False
        pp_release = subprocess.check_output(['vzctl', 'exec2', ct_name, '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_name}.")
            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_name, backup_ct_uuid):
    try:
        ctid=get_ct_id(ct_name)
        subprocess.check_output(['vzmlocal', '--online', '-C', f'{ctid}:{backup_ct_uuid}'], universal_newlines=True)
        subprocess.check_output(['prlctl', 'set', backup_ct_uuid, '--description', ct_name+"_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_name):
    try:
        subprocess.check_output(['vzctl', 'set', ct_name, '--ostemplate', '.vzlinux-7-x86_64', '--distribution', 'vzlinux', '--save'], universal_newlines=True)
        logger.info(f"CT {ct_name} switched successfully to vzlinux distribution.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to convert CT {ct_name}: {e}")
        return False
    return True

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

def revert_backup(ct_name, backup_ct_uuid):
    try:
        subprocess.check_output(['prlctl', 'stop', ct_name], universal_newlines=True)
        subprocess.check_output(['prlctl', 'set', ct_name, "--name", ct_name+"-convert-failed"], universal_newlines=True)
        subprocess.check_output(['prlctl', 'set', backup_ct_uuid, "--name", ct_name], universal_newlines=True)
        subprocess.check_output(['prlctl', 'start', backup_ct_uuid], universal_newlines=True)
        logger.info(f"CT {ct_name} reverted to initial state from backup and started,  resulting container is left under the '{ct_name}-convert-failed' name for investigation")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to revert CT state {ct_name}: {e}")
        return False
    return True

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

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

def run_internal_commands(ct_name):    
    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_name, 'bash', '-c', yum_shell_command], universal_newlines=True)
        logger.info(f"Yum shell commands executed successfully in CT {ct_name}.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to run yum shell commands in CT {ct_name}: {e}")
        return False
    return True

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

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

def check_needs_restarting(ct_name):
    try:
        # Run the needs-restarting command and capture the output and the exit code
        result = subprocess.run(['vzctl', 'exec3', ct_name, '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_name}:\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_name}.")
            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_name} due to core updates.")
                return True
        else:
            logger.warning(f"No restart is necessary for CT {ct_name}.")
            return False
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to check if restart is needed for CT {ct_name}: {e}")
        return False

def validate_name(ct):
    name=None
    try:
        output = subprocess.check_output(['vzlist', '-j', '-H', '-n', ct], universal_newlines=True)
        name= json.loads(output)[0]["name"]
    except:
        try:
            output = subprocess.check_output(['prlctl', 'list', '-j', '-i', ct], universal_newlines=True)
            name= json.loads(output)[0]["Name"]
        except:
            pass
    if name:
        logger.info(f"CT name is {name}")
    else:
        logger.error(f"Invalid CT identificator {ct}")
    return name

def main(ct_name):
    if not check_vzlist(ct_name):
        sys.exit(1)
    
    if not check_config_link(ct_name):
        logger.error(f"Unable to detect {ct_name} config consistency")
        sys.exit(1)

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

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

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

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

    run_distro_sync(ct_name)

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

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


    if check_needs_restarting(ct_name):
        restart_ct(ct_name)

    logger.info(f"All operations completed successfully for CT {ct_name} 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_name | CT_NAME>")
        sys.exit(1)

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

    ct_name=validate_name(ct)
    if ct_name:
        main(ct_name)
    else:
        logger.error(f"Unable to find {f}")
        sys.exit(255)

