# -*- coding: utf-8 -*-
"""
Control-M SaaS migration export utility.

Runs per-component exports, optionally encrypts each with AES-256-GCM,
and bundles everything into a single CTMSExportForSaaS.tar.gz archive.
"""

import argparse
import os
import sys

from ctmsaascommon import (
    Cipher,
    COMPONENT_REGISTRY,
    FAILURE,
    FINAL_ARCHIVE_NAME,
    FLAG_KEY_GLOBAL,
    FLAG_KEY_PBKDF2,
    SUCCESS,
    SUPPORTED_COMPONENTS,
    StepTracker,
    cleanup_files,
    cleanup_work_dir,
    compress_to_targz,
    derive_aes_key,
    encrypt_file_aes_gcm,
    get_server_base_dir,
    load_global_key_bytes,
    logger,
    make_work_dir,
    prompt_password,
    resolve_pack_path,
    setup_logging,
)


def _count_steps(do_encrypt, do_delete_orig, num_components):
    """Compute total progress steps based on runtime options."""
    count = 1                # resolve paths
    count += num_components  # export per component
    if do_encrypt:
        count += 1           # load global key
        count += 1           # encrypt files
        if do_delete_orig:
            count += 1       # remove pre-encryption raw files
    count += 1               # compress to archive
    count += 1               # done
    return count


def _do_export(args, components):
    """Export flow: run component exports, optionally encrypt, bundle into tar.gz."""
    do_encrypt = not args.no_encrypt
    do_delete_orig = do_encrypt and not args.no_delete_orig
    keep_orig = bool(args.no_delete_orig)

    if not do_encrypt and args.password:
        print("WARNING: --password is ignored when --no-encrypt is set.", file=sys.stderr)
        logger.warning("--password ignored because --no-encrypt is set")

    if do_encrypt and Cipher is None:
        print(
            "ERROR: The 'cryptography' package is required for encryption. "
            "Install with: pip install cryptography",
            file=sys.stderr,
        )
        return FAILURE

    tracker = StepTracker(_count_steps(do_encrypt, do_delete_orig, len(components)))

    tracker.step("Resolve server home and paths")
    server_base = os.path.normpath(get_server_base_dir())
    if not os.path.isdir(server_base):
        print("ERROR: Server directory does not exist: %s" % server_base, file=sys.stderr)
        return FAILURE
    logger.info("Server base: %s", server_base)
    archive_path = os.path.join(server_base, FINAL_ARCHIVE_NAME)

    work_dir = make_work_dir(server_base)
    result = FAILURE
    try:
        # --- Per-component export ---
        pack_paths = []
        for comp in components:
            tracker.step("Export component: %s" % comp)
            pack_path = resolve_pack_path(work_dir, comp)
            logger.info("[%s] Pack path: %s", comp, pack_path)

            handler = COMPONENT_REGISTRY.get(comp, {}).get("export")
            if not handler:
                print("ERROR: No export handler for component: %s" % comp, file=sys.stderr)
                return FAILURE
            try:
                handler(server_base, pack_path)
            except Exception as e:
                print("ERROR: [%s] Export failed: %s" % (comp, e), file=sys.stderr)
                logger.error("[%s] Export failed", comp, exc_info=True)
                return FAILURE

            if not os.path.isfile(pack_path):
                print("ERROR: Expected export file missing: %s" % pack_path, file=sys.stderr)
                return FAILURE
            logger.info(
                "[%s] Exported: %s (%d bytes)",
                comp, pack_path, os.path.getsize(pack_path),
            )
            pack_paths.append(pack_path)

        # --- Encrypt (optional) ---
        if do_encrypt:
            tracker.step("Load global key and derive AES-256 encryption key")
            try:
                raw_key = load_global_key_bytes(server_base)
            except Exception as e:
                print("ERROR: %s" % e, file=sys.stderr)
                logger.error("Failed to load global key", exc_info=True)
                return FAILURE

            use_password = bool(args.password)
            password_bytes = b""
            if use_password:
                password_bytes = prompt_password()
                if password_bytes is None:
                    print("ERROR: Too many failed password attempts.", file=sys.stderr)
                    return FAILURE

            aes_key = derive_aes_key(raw_key, use_password, password_bytes)
            flag = FLAG_KEY_PBKDF2 if use_password else FLAG_KEY_GLOBAL

            tracker.step("Encrypt component files with AES-256-GCM")
            enc_paths = []
            for pack_path in pack_paths:
                enc_path = pack_path + ".enc"
                try:
                    encrypt_file_aes_gcm(pack_path, enc_path, aes_key, flag)
                    logger.info("Encrypted: %s -> %s", pack_path, enc_path)
                except Exception as e:
                    print("ERROR: Encryption failed for %s: %s" % (pack_path, e), file=sys.stderr)
                    logger.error("Encryption failed for %s", pack_path, exc_info=True)
                    return FAILURE
                enc_paths.append(enc_path)

            if do_delete_orig:
                tracker.step("Remove pre-encryption raw data files")
                cleanup_files(pack_paths)

            tracker.step("Compress encrypted files to %s" % FINAL_ARCHIVE_NAME)
            try:
                compress_to_targz(enc_paths, archive_path)
                cleanup_files(enc_paths)
            except Exception as e:
                print("ERROR: Compression failed: %s" % e, file=sys.stderr)
                logger.error("Compression failed", exc_info=True)
                return FAILURE

        else:
            tracker.step("Compress component files to %s" % FINAL_ARCHIVE_NAME)
            try:
                compress_to_targz(pack_paths, archive_path)
            except Exception as e:
                print("ERROR: Compression failed: %s" % e, file=sys.stderr)
                logger.error("Compression failed", exc_info=True)
                return FAILURE
            if not keep_orig:
                cleanup_files(pack_paths)

        tracker.step("Done")
        print("Output: %s" % archive_path)
        result = SUCCESS
        return result

    finally:
        if result != SUCCESS and os.path.isfile(archive_path):
            try:
                os.remove(archive_path)
                logger.info("Removed partial archive after failure: %s", archive_path)
            except Exception as e:
                logger.warning("Failed to remove partial archive %s: %s", archive_path, e)

        if keep_orig:
            print("Raw data files kept in: %s" % work_dir)
            logger.info("Work directory kept (--no-delete-orig): %s", work_dir)
        else:
            cleanup_work_dir(work_dir)


def main():
    parser = argparse.ArgumentParser(
        prog="ctmsaasexport",
        description="Export Control-M components with optional AES-256-GCM encryption.",
    )
    parser.add_argument(
        "--component",
        action="append",
        type=str.lower,
        choices=SUPPORTED_COMPONENTS + ["all"],
        metavar="COMPONENT",
        help="Component to export (repeatable). Choices: %s, all. "
             "Defaults to 'all' if omitted." % ", ".join(SUPPORTED_COMPONENTS),
    )
    parser.add_argument(
        "--password",
        action="store_true",
        help="Prompt for a password to derive the encryption key via PBKDF2.",
    )
    parser.add_argument(
        "--no-encrypt",
        action="store_true",
        help="Skip encryption; archive contains raw .txt files.",
    )
    parser.add_argument(
        "--no-delete-orig",
        action="store_true",
        help="Keep raw data files (.txt) after encryption/archiving.",
    )
    args = parser.parse_args()

    setup_logging("ctmsaasexport.log")

    if not args.component or "all" in args.component:
        components = list(SUPPORTED_COMPONENTS)
    else:
        components = list(dict.fromkeys(args.component))

    return _do_export(args, components)


if __name__ == "__main__":
    sys.exit(main())
