#!/usr/bin/env python3

import json
import os
import signal
import sys
import subprocess
import argparse
import shutil
from datetime import datetime

###############################################################################
##                                 VARIABLES                                 ##
###############################################################################

AZURE_CONTAINER_NAME = {
    "stage": "$web",
    "prod":  "$web",
}
AZURE_ACCOUNT_SUFFIX =  {
    "stage": "z20.web.core.windows.net",
    "prod":  "z13.web.core.windows.net",
}
STORAGE_ACCOUNT_NAME =  {
    "stage": "<STORAGE_ACCOUNT_NAME_STG>",
    "prod":  "<STORAGE_ACCOUNT_NAME_PROD>",
}
CDN_RESOURCE_GROUP = {
    "stage": "<RESOURCE_GROUP>",
    "prod":  "<RESOURCE_GROUP>",
}
CDN_PROFILE_NAME = {
    "stage": "<CDN_PROFILE_NAME>",
    "prod":  "<CDN_PROFILE_NAME>",
}


###############################################################################
##                             HELPER FUNCTIONS                              ##
###############################################################################

def signal_handler(sig, frame) -> None:
    """ Signal handler for SIGINT (Ctrl+C) """
    print(f"\n{get_timestamp()} Deployment process interrupted. Exiting gracefully...")
    sys.exit(0)

# Register the signal handler for SIGINT
signal.signal(signal.SIGINT, signal_handler)

def get_timestamp() -> str:
    """ Log prefix timestamp in the format: 'YYYY-MM-DD HH:MM:SS' """
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def colorize(text, color) -> str:
    """ Colorize text for terminal output """
    colors = {
        "red":      "\033[91m",
        "green":    "\033[92m",
        "yellow":   "\033[93m",
        "blue":     "\033[94m",
        "purple":   "\033[95m",
        "cyan":     "\033[96m",
        "white":    "\033[97m",
        "reset":    "\033[0m",
    }
    return f"{colors[color]}{text}{colors['reset']}"

def run_command(command, check=True) -> str:
    """ Run a shell command and return the output """
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    if check and result.returncode != 0:
        print(f"{get_timestamp()} Error: {result.stderr.strip()}")
        sys.exit(result.returncode)
    return result.stdout.strip()

###############################################################################
##                              AZURE FUNCTIONS                              ##
###############################################################################

def check_az_login() -> None:
    """ Check if the user is logged in to Azure CLI """
    try:
        run_command("az account show", check=True)
    except SystemExit:
        print(f"{get_timestamp()} You are not logged in to Azure. Please run 'az login' and try again.")
        sys.exit(1)

def cleanup_old_files(realm) -> None:
    """
        Cleanup old files in Azure Blob Storage.
        LEAD Web will accumulate files over time (as some files have unique names based on the build hash).
        This function will remove all files that were not part of the latest build
    """
    print(f"{get_timestamp()} Cleaning up old files in Azure Blob Storage ({realm})...")
    account_name = STORAGE_ACCOUNT_NAME[realm]
    try:
        # Get list of local files
        local_directory = f"./dist/{realm}"
        local_files = set()
        for root, _, files in os.walk(local_directory):
            for file in files:
                # Collect relative paths as they appear in remote
                relative_path = os.path.relpath(os.path.join(root, file), local_directory)
                local_files.add(relative_path.replace("\\", "/"))  # Ensure uniform path separators

        print(f"{get_timestamp()} Found {len(local_files)} files locally.")

        # Get list of remote files
        list_command = (
            f"az storage blob list "
            f"--account-name {account_name} "
            f"--container-name '{AZURE_CONTAINER_NAME[realm]}' "
            f"--query '[].name' "
            f"--output json"
        )
        remote_files = set(json.loads(run_command(list_command)))

        print(f"{get_timestamp()} Found {len(remote_files)} files in remote storage")

        # Determine files to delete
        files_to_delete = remote_files - local_files
        print(f"{get_timestamp()} Found {len(files_to_delete)} files to delete: {files_to_delete}")

        # Delete the files
        for file in files_to_delete:
            delete_command = (
                f"az storage blob delete "
                f"--account-name {account_name} "
                f"--container-name '{AZURE_CONTAINER_NAME[realm]}' "
                f"--name '{file}' "
            )
            run_command(delete_command)
        print(f"{get_timestamp()} {colorize('Cleanup complete', 'green')} - Removed {len(files_to_delete)} files.")

    except Exception as e:
        print(f"{get_timestamp()} {colorize('Error during cleanup', 'red')}: {str(e)}")
        sys.exit(1)

def purge_cache(realm: str, asyncPurge: bool ) -> None:
    """
        Purge the Azure CDN cache
        Since the other steps clean up old resources, it is important that customers are not pointed to non-existent resources
        This command will take several minutes to run, but can return immediately (and continue asyncronhously) with the --no-wait flag
        For sensitive deploys in which it is essential to receive confirmation of cache purge, do not use the --no-wait flag
        In all other circumstances, asyncPurge set to true will speed the deployment process up
    """
    print(f"{get_timestamp()} Purging Azure CDN cache for realm '{realm}'...")
    try:
        additional_args = " --no-wait true " if asyncPurge else ""
        # Purge the CDN cache
        purge_command = (
            f"az afd endpoint purge "
            f"--resource-group '{CDN_RESOURCE_GROUP[realm]}' "
            f"--profile-name '{CDN_PROFILE_NAME[realm]}' "
            f"--content-paths '/*' "
            f"--endpoint-name '{STORAGE_ACCOUNT_NAME[realm]}' "
            f"{additional_args} "
        )
        run_command(purge_command)
        print(f"{get_timestamp()} {colorize('Cache purge complete', 'green')}")
    except Exception as e:
        print(f"{get_timestamp()} {colorize('Error during cache purge', 'red')}: {str(e)}")
        sys.exit(1)

def deploy(realm: str, cleanup_resources: bool, purge_mode: str) -> None:
    """ Deploy the Angular app to Azure Blob Storage """
    uncommitted_changes = False

    if not realm:
        realm = "stage"
    realm = realm.lower().strip()
    print(f"{get_timestamp()} Starting deployment process for realm '{colorize(realm, 'red')}'...\n")

    # Get git details
    current_branch = run_command("git rev-parse --abbrev-ref HEAD")
    commit_hash = run_command("git rev-parse --short HEAD")
    commit_msg = run_command("git log -1 --pretty=%B")
    print(f"{get_timestamp()} Deploying branch: {current_branch} / commit: {commit_hash} / message: {commit_msg}")

    # Warning unstaged changes: press "Y" to proceed
    if run_command("git status --porcelain"):
        print(f"{get_timestamp()} Warning: There are {colorize('UNSTAGED', 'red')} or {colorize('UNTRACKED', 'red')} changes in the working directory that will be part of the upload.")
        uncommitted_changes = True
        proceed = input("Do you want to proceed with deployment? (Y/N): ")
        if proceed.lower() != "y":
            print(f"{get_timestamp()} Deployment {colorize('aborted', 'red')}.")
            sys.exit(1)

    starttime = datetime.now()

    # Build directory setup
    build_dir = f"./dist/{realm}"
    if os.path.exists(build_dir):
        print(f"{get_timestamp()} Removing old build directory: {build_dir}")
        shutil.rmtree(build_dir)

    # Build the Angular app
    print(f"{get_timestamp()} Building Angular app for realm '{realm}'...")
    build_command = f"./node_modules/@angular/cli/bin/ng build --prod --output-path={build_dir} --configuration={realm}"
    run_command(build_command)

    # Verify the build directory exists
    if not os.path.exists(build_dir):
        print(f"{get_timestamp()} Error: Build directory '{build_dir}' does not exist. Build failed.")
        sys.exit(1)

    # Upload to Azure Blob Storage
    print(f"{get_timestamp()} Uploading files to Azure Blob Storage ({realm})...")
    upload_command = (
        f"az storage blob upload-batch "
        f"--account-name {STORAGE_ACCOUNT_NAME[realm]} "
        f"--source {build_dir} "
        f"--destination '{AZURE_CONTAINER_NAME[realm]}' "
        f"--overwrite "
    )
    run_command(upload_command)

    git_hash = f"{commit_hash}" + f"-{realm}" + f"-{datetime.now().strftime('%Y%m%d%H%M%S')}"
    if uncommitted_changes:
        git_hash += "-ALTERED"
    print(f"{get_timestamp()} Tagging deployment with commit hash: {git_hash}")

    # Tag the deployment with the commit hash
    upload_command = (
        f"az storage blob metadata update "
        f"--account-name {STORAGE_ACCOUNT_NAME[realm]} "
        f"--container-name '$web' "
        f"--name 'index.html' "
        f"--metadata version='{git_hash}' "
    )
    run_command(upload_command)

    if cleanup_resources:
        cleanup_old_files(realm)

    end_deploy_time = datetime.now()

    # Invalidate the CDN so it points to the new resources rather than the old resources (which we've deleted)
    if purge_mode == "sync" or purge_mode == "async":
        purge_cache(realm, asyncPurge=(purge_mode == "async"))

    print(f"{get_timestamp()} {colorize('Deployment complete', 'green')} (deploy time taken: {end_deploy_time - starttime})\n")
    print(f"Direct access at: https://{STORAGE_ACCOUNT_NAME[realm]}.{AZURE_ACCOUNT_SUFFIX[realm]}/index.html\n")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Deploy Angular app to Azure Blob Storage.")
    parser.add_argument("--realm", required=False, help="Deployment realm ('stage', 'prod'). Default: 'stage'.")
    parser.add_argument("--cleanup-resources", required=False, help="Cleanup old resources from previous deploy after current deployment succeeds", default=True)

    cdn_group = parser.add_mutually_exclusive_group()
    cdn_group.add_argument("--no-cdn-purge", action="store_true", help="Do not purge CDN (default).")
    cdn_group.add_argument("--purge-cdn-sync", action="store_true", help="Purge CDN synchronously (will be slow).")
    cdn_group.add_argument("--purge-cdn-async", action="store_true", help="Purge CDN asynchronously (returns immediately).")

    args = parser.parse_args()

    if args.purge_cdn_sync:
        purge_mode = "sync"
    elif args.purge_cdn_async:
        purge_mode = "async"
    else:
        purge_mode = "no-purge"  # Default

    # Check Azure CLI login
    check_az_login()

    # Deploy the app
    deploy(realm=args.realm, cleanup_resources=args.cleanup_resources, purge_mode=purge_mode)