import argparse
import json
import re
import time
import uuid
from enum import Enum

import boto3
from mypy_boto3_stepfunctions import SFNClient
from mypy_boto3_stepfunctions.type_defs import StartExecutionOutputTypeDef

from config import (
    AWS_REGION,
    BULK_ASSET_DOWNLOAD_SFN_NAME,
    SFN_ARN,
)
from utils.cli_input import query_yes_no


class AccountType(Enum):
    VENDOR = "vendor"
    SUBACCOUNT = "subaccount"


class AssetTypes(Enum):
    AUDIO = "WAV"
    VIDEO = "VIDEO_MASTER"
    IMAGE = "TIF"


def main() -> None:
    args = get_args()

    export_id = generate_export_id()

    sfn = boto3.client("stepfunctions", region_name=AWS_REGION)
    response = start_sfn_execution(sfn, args, export_id)

    enforce_spreadsheet_entry(response, args, export_id)
    poll_for_completion(sfn, response["executionArn"])


def generate_export_id() -> uuid.UUID:
    generate_new_export = query_yes_no(
        "Do you want to generate a new export ID for this request? "
    )
    if not generate_new_export:
        user_export_id = get_input_export_id()
        # Make mypy happy
        if not user_export_id:
            raise ValueError("Export ID cannot be None.")
        return user_export_id
    return uuid.uuid4()


def get_input_export_id() -> uuid.UUID | None:
    while True:
        user_export_id = input("Please enter the export UUID for this request: ")
        try:
            export_id = uuid.UUID(user_export_id, version=4)
        except ValueError:
            print("Invalid export ID. Please enter a valid UUID4 and try again.")
            continue

        return export_id


def start_sfn_execution(
    sfn: SFNClient, args: argparse.Namespace, export_id: uuid.UUID
) -> StartExecutionOutputTypeDef:
    sfn_input = {
        "account_id": args.account_id,
        "account_type": args.account_type,
        "asset_types": [x.upper() for x in args.asset_types],
        "export_id": str(export_id),
    }
    print(
        f"Kicking off {BULK_ASSET_DOWNLOAD_SFN_NAME} "
        f"execution for export {export_id}...\n"
    )
    response = sfn.start_execution(stateMachineArn=SFN_ARN, input=json.dumps(sfn_input))
    time.sleep(2)
    print("Execution started!\n")
    return response


def enforce_spreadsheet_entry(
    response: StartExecutionOutputTypeDef,
    args: argparse.Namespace,
    export_id: uuid.UUID,
) -> None:
    start_date = (
        response["startDate"].isoformat(timespec="milliseconds").replace("+00:00", "Z")
    )
    time.sleep(2)
    print(
        "Please fill in the export spreadsheet with the "
        "following details while the execution runs...\n"
    )
    time.sleep(2)
    print(f"EXPORT_ID: {export_id}")
    print(f"ACCOUNT_ID: {args.account_id}")
    print(f"ACCOUNT_TYPE: {args.account_type}")
    print(f"ASSET_TYPES: {', '.join(args.asset_types)}")
    print(f"EXECUTION_NAME: {response['executionArn'].split(':')[-1]}")
    print(f"EXPORT_START_DATE: {start_date}\n")
    print(f"USERNAME: {args.username}\n")
    print(f"USER_EMAIL: {args.user_email}\n")
    print(f"JIRA_TICKET: {args.jira_ticket.upper()}\n")

    while True:
        spreadsheet_entered = query_yes_no(
            "Have you filled in the export spreadsheet with the details above? "
        )
        if spreadsheet_entered:
            print("Thank you!\n")
            time.sleep(2)
            break
        print("Please fill in the export spreadsheet before proceeding.\n")


def poll_for_completion(sfn: SFNClient, execution_arn: str) -> None:
    sleep_time_seconds = 30
    attempts = 0
    max_attempts = 5  # After ~15 minutes, give up and suggest checking console
    while True:
        print("Checking execution status...\n")
        time.sleep(2)
        response = sfn.describe_execution(executionArn=execution_arn)
        status = response["status"]
        if status in ["SUCCEEDED", "FAILED", "TIMED_OUT", "ABORTED"]:
            if status == "SUCCEEDED":
                print(f"Execution {status.lower()}! 🎉 🪩\n")
                end_date = (
                    response["stopDate"]
                    .isoformat(timespec="milliseconds")
                    .replace("+00:00", "Z")
                )
                time.sleep(1)
                print("***Please enter the END_DATE into the export spreadsheet...***")
                print(f"EXPORT_END_DATE: {end_date}\n")
                time.sleep(1)

                print("Thank you for using the Bulk Asset Downloader! Goodbye! 👋")
                exit(0)
            else:
                print(
                    f"Status: {status}\n"
                    f"\nExecution {execution_arn} did not complete successfully.\n"
                    f"Please check the AWS console for more details."
                )
                exit(1)
        if attempts > max_attempts:
            print(
                "This appears to be a long running task, check the "
                "AWS console later...\n"
                "Don't forget to enter the END_DATE into the export "
                "spreadsheet when it completes.\n"
                "Goodbye! 👋"
            )
            exit(0)
        print(f"Execution is still running... (status: {status})")
        time.sleep(sleep_time_seconds)
        sleep_time_seconds *= 2  # Exponential backoff
        attempts += 1


def get_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "account_id",
        type=int,
        help="OA Label ID.",
    )
    parser.add_argument(
        "account_type",
        type=str,
        choices=[x.value for x in AccountType],
        help="Type of account: vendor or subaccount.",
    )
    parser.add_argument(
        "asset_types",
        type=lambda s: s.split(","),
        help="Comma-separated list of asset types to download: "
        "(e.g. wav,video_master,tif).",
    )
    parser.add_argument(
        "username",
        type=str,
        help="Your username (e.g. mclark).",
    )
    parser.add_argument(
        "user_email",
        type=str,
        help="Your user email (e.g. mikayla.clark@sonymusic-pde.com).",
    )
    parser.add_argument(
        "jira_ticket",
        type=str,
        help="JIRA Ticket associated with the asset request (e.g. cdam-1234).",
    )
    args = parser.parse_args()

    # Validate JIRA ticket format
    if not re.match(r"^[a-zA-Z]+-\d+$", args.jira_ticket):
        print(f"Invalid JIRA ticket ID: {args.jira_ticket}\n")
        print("Please enter a valid JIRA ticket ID (e.g. CDAM-1234).\n")
        exit(1)
    # Validate asset types
    valid_asset_types = [item.value for item in AssetTypes]
    input_asset_types = [x.upper() for x in args.asset_types]
    for asset_type in input_asset_types:
        if asset_type not in valid_asset_types:
            print(
                f"Invalid asset type: {asset_type}. "
                f"Valid asset types are: {','.join(valid_asset_types)}"
            )
            exit(1)
    return args


if __name__ == "__main__":
    main()