#!/bin/python3
# To download and install Docker Language and GPU Images from Ai-media AWS S3 bucket
import argparse
import logging
import subprocess
import os

from datetime import datetime
from pathlib import Path

LOG_NAME = "UpdateDockerImages"
LOG_LEVEL = logging.DEBUG
LOG_FORMAT = "%(asctime)s - %(levelname)-5s - %(message)s"
DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
LOG_FILENAME = f"update_docker_images_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
LOG_PATH = Path("/home/eeg") / LOG_FILENAME
logger = logging.getLogger(LOG_NAME)
logger.setLevel(LOG_LEVEL)
CONSOLE_HANDLER = logging.StreamHandler()
FILE_HANDLER = logging.FileHandler(LOG_PATH)
FORMATTER = logging.Formatter(LOG_FORMAT, datefmt=DATE_FORMAT)
CONSOLE_HANDLER.setFormatter(FORMATTER)
FILE_HANDLER.setFormatter(FORMATTER)
logger.addHandler(CONSOLE_HANDLER)
logger.addHandler(FILE_HANDLER)

# BASE URL AND VERSION
BASE_URL = "https://eeg-customer-deliverables.s3.us-east-1.amazonaws.com/lexi_local/Docker_Images"
LANG_PATH = 'Languages'
GPU_PATH = 'GPU'

DEFAULT_LANG_IMAGES = ['english', 'french', 'german', 'spanish']
DEFAULT_GPU_IMAGES = ['recipe1', 'recipe2', 'recipe3', 'recipe4']

images_successful = []
images_error = []

DESCRIPTION = """
Script for update the Docker Images for Lexi Local/DR

Global Options:
  -y, --yes               Skip confirmation prompts (automated mode)
  -v, --version           Docker image version
  -l, --langs             List of language images to update (comma-separated)
  -nl, --no-langs         Do not update language images
  -g, --gpu               List of GPU images to update (comma-separated)
  -ng, --no-gpu           Do not update GPU images
  --print-config          Print the current configuration

Notes:
- Use -y/--yes to skip confirmation prompts.
- Use --print-config to display the current configuration.
"""


###########################################
# CONFIGURATION
###########################################
class Config:
    AUTO_CONFIRM = False
    LANG_IMAGES = [DEFAULT_LANG_IMAGES]
    GPU_IMAGES = [DEFAULT_GPU_IMAGES]
    DOCKER_IMAGE_VERSION = None

    @staticmethod
    def print_all():
        logger.info("=== Current Configuration ===")
        for attr in dir(Config):
            if attr.isupper():
                logger.info(f"{attr}: {getattr(Config, attr)}")
        logger.info("============================")

    @staticmethod
    def update_global_vars(args, parser):
        # Configuration options
        Config.AUTO_CONFIRM = args.yes
        Config.DOCKER_IMAGE_VERSION = args.version

        # Images
        Config.LANG_IMAGES = args.langs.split(',') if args.langs else DEFAULT_LANG_IMAGES if not args.no_langs else []
        Config.GPU_IMAGES = args.gpu.split(',') if args.gpu else DEFAULT_GPU_IMAGES if not args.no_gpu else []


###########################################
# HELPERS
###########################################
def prompt_yes_no(text, default=True):
    if Config.AUTO_CONFIRM:
        return default

    options = "(Y/n)" if default else "(y/N)"

    while True:
        res = input(f"{text} {options}: ").strip().lower()

        if res in ('y', 'yes'):
            return True
        elif res in ('n', 'no'):
            return False
        elif res == 'x':
            logger.error("❌ Exiting by user request.")
            exit(0)
        elif res == '':
            return default
        else:
            logger.error(f"⚠️ Invalid response: '{res}'. Please enter Y, N, or X to exit.")


###########################################
# SERVICES OPERATIONS
###########################################
def restart_service(service_name):
    try:
        logger.info(f"Restarting service: {service_name}")
        subprocess.run(['systemctl', 'restart', service_name], check=True)
        logger.info(f"✅ Service restarted: {service_name}")
        return True
    except Exception as e:
        logger.error(f"❌ Error restarting service {service_name}: {e}")
        return False


###########################################
# IMAGE DOWNLOAD HELPERS
###########################################
def download_image(tar_filename, url):
    """
    Download the Docker image tar file from the specified URL.

    Args:
        tar_filename (str): The name of the tar file to save.
        url (str): The URL to download the tar file from.
    Returns:
        bool: True if the download was successful, False otherwise.
    """

    logger.info(f'Step 1/3: Downloading {tar_filename} from {url}')
    try:
        result = subprocess.run(["wget", url])
        return result.returncode == 0
    except subprocess.CalledProcessError as e:
        logger.error(f"[!] Error downloading {tar_filename}: {e}")
        return False


def load_docker_image(tar_filename):
    """
    Load the Docker image from the specified tar file.
    Args:
        tar_filename (str): The name of the tar file to load.
    Returns:
        bool: True if the image was loaded successfully, False otherwise.
    """

    logger.info(f"Step 2/3: Loading {tar_filename} into Docker")

    try:
        result = subprocess.run(["docker", "image", "load", "-i", tar_filename])
        return result.returncode == 0
    except subprocess.CalledProcessError as e:
        logger.error(f"[!] Error loading {tar_filename} into Docker: {e}")
        return False


def delete_compressed_file(tar_filename):
    """
    Delete the specified tar file after loading it into Docker.
    Args:
        tar_filename (str): The name of the tar file to delete.
    Returns:
        bool: True if the file was deleted successfully, False otherwise.
    """

    logger.info(f"Step 3/3: Deleting {tar_filename}")

    try:
        os.remove(tar_filename)
        return True
    except OSError as e:
        logger.error(f"[!] Error removing {tar_filename}: {e}")
        return False


def process_image_prompt(image):
    text = f"Do you want to update  the image '{image}'?"
    return prompt_yes_no(text, default=True)


def process_image(image, type_image):
    """
    Process a single Docker image by downloading, loading, and deleting the tar file.
    Args:
        image (str): The name of the Docker image to process.
        type_image (str): The type of image (e.g., 'Languages' or 'GPU').
    Returns:
        bool: True if the image was processed successfully, False otherwise.
    """

    if not process_image_prompt(image):
        logger.warning(f"❌ Skipping image: {image}")
        return False

    tar_filename = f"{image}_{Config.DOCKER_IMAGE_VERSION}.tar"
    url = f"{BASE_URL}/{type_image}/{Config.DOCKER_IMAGE_VERSION}/{tar_filename}"

    # Download the image
    if not download_image(tar_filename, url):
        logger.error(f"[!] Download failed for {tar_filename}. Skipping to next image.")
        images_error.append([image, 'Not found'])

        return False
    logger.info('Image downloaded!')

    # Load the image into Docker
    if not load_docker_image(tar_filename):
        logger.error(f"[!] Docker load failed for {tar_filename}. Skipping deletion.")
        images_error.append([image, 'Not loaded'])

        return False
    logger.info('Image loaded into Docker!')

    # Delete the local tar file
    if not delete_compressed_file(tar_filename):
        logger.error(f"[!] Failed to delete {tar_filename} after loading into Docker.")
        images_error.append([image, 'Not deleted'])

        return False
    logger.info('Local tar file deleted!')

    return True


def process_images(images, type_image):
    for image in images:
        logger.info(f"\nProcessing image: {image}")

        if process_image(image, type_image):
            images_successful.append(image)

        logger.info('-' * 50 + '\n')


###########################################
# IMAGE DOWNLOAD FUNCTIONS
###########################################
def update_image_prompt(image_type):
    text = f"Do you want to update {image_type} images'?"
    return prompt_yes_no(text, default=True)


def download_transcription_images():
    logger.info("🔄  Download Transcription images.")

    if len(Config.LANG_IMAGES) == 0:
        logger.warning("⚠️ No language images to download.")
        return

    if not update_image_prompt('Transcription'):
        logger.warning("❌ Skipping transcription images update.")
        return

    process_images(Config.LANG_IMAGES, LANG_PATH)


def download_gpu_images():
    logger.info("🔄  Download GPU images.")

    if len(Config.GPU_IMAGES) == 0:
        logger.warning("⚠️ No GPU images to download.")
        return

    if not update_image_prompt('GPU'):
        logger.warning("❌ Skipping GPU images update.")
        return

    process_images(Config.GPU_IMAGES, GPU_PATH)


###########################################
# MAIN
###########################################
def add_arguments(parser):
    argument_groups = {
        'General': [
            (("--yes", "-y"), "store_true", "Skip confirmation prompt (automated mode)"),
            (("--version", "-v"), "store", "Docker image version"),
            (("--langs", "-l"), "store", "List of language images to update (comma-separated)"),
            (("--no-langs", "-nl"), "store_true", "Do not update language images"),
            (("--gpu", "-g"), "store", "List of GPU images to update (comma-separated)"),
            (("--no-gpu", "-ng"), "store_true", "Do not update GPU images"),
            (("--print-config",), "store_true", "Print the current configuration"),
        ],
    }

    for group, args_list in argument_groups.items():
        group_parser = parser.add_argument_group(group)
        for arg_names, action, help_text in args_list:
            is_required = "--version" in arg_names
            group_parser.add_argument(*arg_names, action=action, help=help_text, required=is_required)


def print_config_if_requested(args):
    if args.print_config:
        Config.print_all()
        return True

    return False


def download_all_images():
    download_transcription_images()
    download_gpu_images()


def summarize_results():
    logger.info("=== Summary ===")
    for category, items in [('Errors', images_error), ('Successes', images_successful)]:
        if items:
            logger.info(f"✅ {category}:")
            for item in items:
                logger.info(f"  - {item}")
        else:
            logger.info(f"⚠️ No {category.lower()} encountered.")
    logger.info("====================")


def restart_services():
    """
    Restart the necessary services after updating Docker images.
    """

    logger.info("🔄 Restarting services.")

    services = ['usage_monitor', 'lexi_wsgi',]

    for service in services:
        restart_service(service)

    logger.info("Services restarted successfully.")


def main():
    parser = argparse.ArgumentParser(description=DESCRIPTION, formatter_class=argparse.RawDescriptionHelpFormatter)
    add_arguments(parser)

    args = parser.parse_args()

    Config.update_global_vars(args, parser)

    if print_config_if_requested(args):
        return

    logger.info("==== Script to download and install Docker Language and GPU Images from Ai-media AWS S3 bucket ====")
    logger.info(f'Version to download: {Config.DOCKER_IMAGE_VERSION}\n')

    # Download all images
    download_all_images()

    logger.info("=== Update finished ===")

    # Print summary of results (errors, successes, skipped)
    summarize_results()

    # Restart services if images were successfully updated
    if len(images_successful) > 0:
        restart_services()


if __name__ == "__main__":
    main()
