# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

"""
All the `deadline bundle` commands.
"""

from __future__ import annotations

import json
import logging
import sys
import re
from typing import Any, Optional
import tempfile
import shutil
import os
from dataclasses import fields

import click
from botocore.exceptions import ClientError

from ... import api
from ...config import config_file
from ...dataclasses import SubmitterInfo
from ....job_attachments.exceptions import (
    AssetSyncError,
    AssetSyncCancelledError,
    MisconfiguredInputsError,
)
from ....job_attachments.models import JobAttachmentsFileSystem

from ...exceptions import DeadlineOperationError, CreateJobWaiterCanceled
from .._common import (
    _apply_cli_options_to_config,
    _handle_error,
    _ProgressBarCallbackManager,
    _parse_multi_format_parameters,
    _suggest_resources_on_client_error,
)
from .._main import deadline as main
from ._sigint_handler import SigIntHandler

logger = logging.getLogger(__name__)

# Set up the signal handler for handling Ctrl + C interruptions.
sigint_handler = SigIntHandler()


@main.group(name="bundle")
@_handle_error
def cli_bundle():
    """
    Commands to work with Open Job Description [job bundles]. Use these commands to
    submit jobs to run on a Deadline Cloud queue.

    [job bundles]: https://docs.aws.amazon.com/deadline-cloud/latest/developerguide/build-job-bundle.html
    """


# Latin alphanumeric, starting with a letter
_openjd_identifier_regex = r"(?-m:^[A-Za-z_][A-Za-z0-9_]*\Z)"


def validate_parameters(ctx, param, value):
    """
    Validate provided --parameter values, ensuring that they are in the format "ParamName=Value", and convert them to a dict with the
    following format:
        [{"name": "<name>", "value": "<value>"}, ...]
    """
    parameters_split = []
    for parameter in value:
        regex_match = re.match("([^=]+)=(.*)", parameter)
        if not regex_match:
            raise click.BadParameter(
                f'Parameters must be provided in the format "ParamName=Value". Invalid parameter: {parameter}'
            )

        if not re.match(_openjd_identifier_regex, regex_match[1]):
            raise click.BadParameter(
                f"Parameter names must be alphanumeric Open Job Description identifiers. Invalid parameter name: {regex_match[1]}"
            )

        parameters_split.append({"name": regex_match[1], "value": regex_match[2]})

    return parameters_split


def _validate_submitter_info(ctx, param, values):
    """
    Validate provided --submitter-info value and convert to SubmitterInfo object.

    Supports three input formats that can be mixed:
    - Key=value pairs: --submitter-info submitter_name=MyApp --submitter-info host_application_name=Maya
    - Inline JSON strings: --submitter-info '{"submitter_name": "MyApp", "additional_info": {"custom": "data"}}'
    - File paths (JSON or YAML): --submitter-info file://path/to/submitter.json

    All keys must be valid SubmitterInfo fields. Unknown keys will raise an error.
    """
    if not values:
        return None

    # Get valid field names from SubmitterInfo dataclass
    valid_fields = {field.name for field in fields(SubmitterInfo)}

    info_dict = _parse_multi_format_parameters(list(values))

    # Validate all keys
    for key in info_dict.keys():
        if key not in valid_fields:
            raise click.BadParameter(
                f"Unknown field '{key}'. Valid fields are: {', '.join(sorted(valid_fields))}"
            )

    # Ensure submitter_name is provided as a required field
    if "submitter_name" not in info_dict:
        raise click.BadParameter(
            "submitter_name is required when using --submitter-info. "
            "Example: --submitter-info submitter_name=MyApp"
        )

    try:
        return SubmitterInfo(**info_dict)
    except TypeError as e:
        raise click.BadParameter(f"Failed to create SubmitterInfo: {e}") from e


def _interactive_confirmation_prompt(message: str, default_response: bool) -> bool:
    """
    Callback to decide if submission should continue or be canceled. Returns True to continue, False to cancel.

    Args:
        warning_message (str): The warning message to display.
        default_response (bool): The default to present as the response (True to continue, False to cancel).
    """
    return click.confirm(
        message,
        default=default_response,
    )


@cli_bundle.command(name="submit")
@click.option(
    "-p",
    "--parameter",
    multiple=True,
    callback=validate_parameters,
    help=(
        "The values for the job template's parameters. Can be provided as key-value pairs, inline JSON strings, "
        "or as paths to a JSON or YAML document. Later values for repeated parameter names take precedence. "
        'Examples: --parameter MyParam=5 -p file://parameter_file.json -p \'{"OtherParam": "10"}\''
    ),
)
@click.option("--profile", help="The AWS profile to use.")
@click.option("--farm-id", help="The farm to use.")
@click.option("--queue-id", help="The queue to use.")
@click.option("--storage-profile-id", help="The storage profile to use.")
@click.option("--name", help="The job name to use in place of the one in the job bundle.")
@click.option(
    "--priority",
    type=int,
    default=50,
    help="The priority of the job. Jobs with a higher priority run first.",
)
@click.option(
    "--max-failed-tasks-count",
    type=int,
    help="The maximum number of failed tasks before the job is marked as failed.",
)
@click.option(
    "--max-retries-per-task",
    type=int,
    help="The maximum number of times to retry a task before it is marked as failed.",
)
@click.option(
    "--max-worker-count",
    type=int,
    help="The max worker count of the job.",
)
@click.option(
    "--target-task-run-status",
    type=click.Choice(["READY", "SUSPENDED"], case_sensitive=False),
    help="The target task run status for the job. READY means tasks will start immediately, "
    "SUSPENDED means tasks will be created but not start until manually resumed.",
)
@click.option(
    "--job-attachments-file-system",
    help="The method workers use to access job attachments. "
    "COPIED means to copy files to the worker and VIRTUAL means to load "
    "files as needed from a virtual file system. If VIRTUAL is selected "
    "but not supported by a worker, it will fallback to COPIED.",
    type=click.Choice([e.value for e in JobAttachmentsFileSystem]),
)
@click.option(
    "--yes",
    is_flag=True,
    help="Automatically accept any confirmation prompts",
)
@click.option(
    "--require-paths-exist",
    is_flag=True,
    help="Return an error if any input files are missing.",
)
@click.option(
    "--submitter-name",
    type=click.STRING,
    help="Name of the application submitting the bundle.",
)
@click.option(
    "--known-asset-path",
    multiple=True,
    help="Path that should not generate warnings when outside storage profile locations. "
    "Can be specified multiple times for different paths.",
)
@click.option(
    "--save-debug-snapshot",
    help="EXPERIMENTAL - Instead of submitting the job, generate a debug snapshot as a directory or a zip file if the extension is .zip."
    " It includes the job attachments and parameters for creating the job."
    " You can later run the bash script in the snapshot to submit the job using AWS CLI commands.",
)
@click.option(
    "--force-s3-check/--no-force-s3-check",
    default=None,
    help="Force verification that job attachments exist in S3 before skipping upload. "
    "Use when S3 bucket contents may be out of sync with local caches. "
    "Overrides the 'settings.force_s3_check' config setting.",
)
@click.argument("job_bundle_dir")
@_handle_error
def bundle_submit(
    job_bundle_dir,
    job_attachments_file_system,
    parameter,
    known_asset_path,
    name,
    priority,
    max_failed_tasks_count,
    max_retries_per_task,
    max_worker_count,
    target_task_run_status,
    require_paths_exist,
    submitter_name,
    save_debug_snapshot,
    force_s3_check,
    **args,
):
    """
    Submits an Open Job Description [job bundle] to a
    [Deadline Cloud queue]. You can provide options
    to set parameter values, the job name, priority, and more.

    [job bundle]: https://docs.aws.amazon.com/deadline-cloud/latest/developerguide/build-job-bundle.html
    [Deadline Cloud queue]: https://docs.aws.amazon.com/deadline-cloud/latest/userguide/queues.html
    """
    # Apply the CLI args to the config
    config = _apply_cli_options_to_config(required_options={"farm_id", "queue_id"}, **args)

    # Resolve force_s3_check: CLI flag takes precedence, otherwise use config setting
    if force_s3_check is None:
        force_s3_check = config_file.str2bool(
            config_file.get_setting("settings.force_s3_check", config=config)
        )

    hash_callback_manager = _ProgressBarCallbackManager(length=100, label="Hashing Attachments")
    upload_callback_manager = _ProgressBarCallbackManager(length=100, label="Uploading Attachments")

    def _check_create_job_wait_canceled() -> bool:
        return sigint_handler.continue_operation

    try:
        snapshot_tmpdir = None
        if save_debug_snapshot:
            save_debug_snapshot = os.path.abspath(save_debug_snapshot)

            # If the debug snapshot is to a zip file, first put it in a temporary directory
            if save_debug_snapshot.endswith(".zip"):
                snapshot_tmpdir = tempfile.TemporaryDirectory()

        job_id = api.create_job_from_job_bundle(
            job_bundle_dir=job_bundle_dir,
            job_parameters=parameter,
            name=name,
            job_attachments_file_system=job_attachments_file_system,
            config=config,
            priority=priority,
            max_failed_tasks_count=max_failed_tasks_count,
            max_retries_per_task=max_retries_per_task,
            max_worker_count=max_worker_count,
            target_task_run_status=target_task_run_status,
            hashing_progress_callback=hash_callback_manager.callback,
            upload_progress_callback=upload_callback_manager.callback,
            create_job_result_callback=_check_create_job_wait_canceled,
            print_function_callback=click.echo,
            interactive_confirmation_callback=_interactive_confirmation_prompt,
            require_paths_exist=require_paths_exist,
            submitter_name=submitter_name or "CLI",
            known_asset_paths=known_asset_path,
            debug_snapshot_dir=snapshot_tmpdir.name if snapshot_tmpdir else save_debug_snapshot,
            force_s3_check=force_s3_check,
        )

        if snapshot_tmpdir:
            # Put the snapshot in a zip file
            os.makedirs(os.path.dirname(save_debug_snapshot), exist_ok=True)
            shutil.make_archive(save_debug_snapshot, "zip", snapshot_tmpdir.name)

        if save_debug_snapshot:
            click.echo("Saved job debug snapshot:")
            click.echo(f"    {save_debug_snapshot}")

        # Check Whether the CLI options are modifying any of the default settings that affect
        # the job id. If not, we'll save the job id submitted as the default job id.
        # If a job snapshot directory was provided, the job_id will be None.
        if (
            args.get("profile") is None
            and args.get("farm_id") is None
            and args.get("queue_id") is None
            and args.get("storage_profile_id") is None
            and job_id
        ):
            config_file.set_setting("defaults.job_id", job_id)

    except AssetSyncCancelledError as exc:
        if sigint_handler.continue_operation:
            raise DeadlineOperationError(f"Job submission unexpectedly canceled:\n{exc}") from exc
        else:
            click.echo("Job submission canceled.")
            sys.exit(1)
    except AssetSyncError as exc:
        raise DeadlineOperationError(f"Failed to upload job attachments:\n{exc}") from exc
    except CreateJobWaiterCanceled as exc:
        if sigint_handler.continue_operation:
            raise DeadlineOperationError(
                f"Unexpectedly canceled during wait for final status of CreateJob:\n{exc}"
            ) from exc
        else:
            click.echo("Canceled waiting for final status of CreateJob.")
            sys.exit(1)
    except ClientError as exc:
        suggestion = _suggest_resources_on_client_error(
            exc,
            farm_id=config_file.get_setting("defaults.farm_id", config=config),
            queue_id=config_file.get_setting("defaults.queue_id", config=config),
            config=config,
        )
        raise DeadlineOperationError(
            f"Failed to submit the job bundle to AWS Deadline Cloud:\n{exc}{suggestion}"
        ) from exc
    except MisconfiguredInputsError as exc:
        click.echo(str(exc))
        click.echo("Job submission canceled.")
        sys.exit(1)
    except Exception as exc:
        api.get_deadline_cloud_library_telemetry_client().record_error(
            event_details={"exception_scope": "on_submit"},
            exception_type=str(type(exc)),
        )
        raise
    finally:
        if snapshot_tmpdir:
            snapshot_tmpdir.cleanup()


@cli_bundle.command(name="gui-submit")
@click.option(
    "-p",
    "--parameter",
    multiple=True,
    callback=validate_parameters,
    help=(
        "Initial values in the GUI for the job template's parameters. Can be provided as key-value pairs, inline JSON strings, "
        "or as paths to a JSON or YAML document. Later values for repeated parameter names take precedence. "
        'Examples: --parameter MyParam=5 -p file://parameter_file.json -p \'{"OtherParam": "10"}\''
    ),
)
@click.argument("job_bundle_dir", required=False)
@click.option(
    "--browse",
    is_flag=True,
    help="Opens a folder browser to select a bundle.",
)
@click.option(
    "--install-gui",
    is_flag=True,
    help="Installs GUI dependencies if they are not installed already",
)
@click.option(
    "--submitter-name",
    help="[DEPRECATED] Use --submitter-info submitter_name=<name> instead. Name of the application submitting the bundle. If a name is specified, the GUI will automatically close after submitting the job.",
)
@click.option(
    "--output",
    type=click.Choice(
        ["verbose", "json"],
        case_sensitive=False,
    ),
    default="verbose",
    help="Specifies the output format of the messages printed to stdout.\n"
    "VERBOSE: Displays messages in a human-readable text format.\n"
    "JSON: Displays messages in JSON line format, so that the info can be easily "
    "parsed/consumed by custom scripts.",
)
@click.option(
    "--known-asset-path",
    multiple=True,
    help="Path that should not generate warnings when outside storage profile locations. "
    "Can be specified multiple times for different paths.",
)
@click.option(
    "--submitter-info",
    multiple=True,
    callback=_validate_submitter_info,
    help="Submitter and environment information. Supports key=value pairs, inline JSON strings, "
    "and file paths (JSON or YAML). Later values for repeated fields take precedence. "
    "Examples: --submitter-info submitter_name=MyApp --submitter-info host_application_name=Maya "
    'OR --submitter-info \'{"submitter_name": "MyApp", "additional_info": {"render_engine": "Cycles"}}\' '
    "OR --submitter-info file://path/to/submitter.json",
)
@_handle_error
def bundle_gui_submit(
    parameter,
    job_bundle_dir,
    browse,
    output,
    install_gui,
    known_asset_path,
    submitter_name,
    submitter_info,
    **args,
):
    """
    Opens a GUI to submit an Open Job Description [job bundle] to a
    [Deadline Cloud queue]. You can provide options
    to set the initial parameter values shown in the GUI.

    [job bundle]: https://docs.aws.amazon.com/deadline-cloud/latest/developerguide/build-job-bundle.html
    [Deadline Cloud queue]: https://docs.aws.amazon.com/deadline-cloud/latest/userguide/queues.html
    """

    if submitter_name:
        click.echo(
            click.style(
                "DeprecationWarning: The option --submitter-name is deprecated. Use --submitter-info instead.",
                fg="red",
            ),
            err=True,
        )
        if submitter_info:
            # --submitter-name takes precedence if we already have submitter_info provided
            submitter_info.submitter_name = submitter_name
        else:
            submitter_info = SubmitterInfo(submitter_name=submitter_name)

    from ...ui import gui_context_for_cli
    from ...ui._utils import tr

    with gui_context_for_cli(automatically_install_dependencies=install_gui) as app:
        from ...ui.job_bundle_submitter import show_job_bundle_submitter

        if not job_bundle_dir and not browse:
            raise DeadlineOperationError(
                tr(
                    "Specify a job bundle directory or run the bundle command with the --browse flag"
                )
            )
        output = output.lower()

        submitter = show_job_bundle_submitter(
            input_job_bundle_dir=job_bundle_dir,
            browse=browse,
            submitter_info=submitter_info,
            known_asset_paths=known_asset_path,
            job_parameters=parameter,
        )

        if not submitter:
            return

        submitter.show()

        app.exec()

        _print_response(
            output=output,
            job_bundle_dir=job_bundle_dir,
            job_history_bundle_dir=submitter.job_history_bundle_dir,
            job_id=submitter.job_id,
        )


def _print_response(
    output: str,
    job_bundle_dir: str,
    job_history_bundle_dir: Optional[str],
    job_id: Optional[str],
):
    if output == "json":
        if job_id:
            response: dict[str, Any] = {
                "status": "SUBMITTED",
                "jobId": job_id,
                "jobHistoryBundleDirectory": job_history_bundle_dir,
            }
            click.echo(json.dumps(response))
        else:
            click.echo(json.dumps({"status": "CANCELED"}))
    else:
        if job_id:
            click.echo("Submitted job bundle:")
            click.echo(f"   {job_bundle_dir}")
            click.echo(f"Job ID: {job_id}")
        else:
            click.echo("Job submission canceled.")
