import typer
from rich import print
from typing_extensions import Annotated

from airflow_tools.connectors.mwaa import MWAAConnector
from airflow_tools.connectors.s3 import S3Connector

app: typer.Typer = typer.Typer()


@app.command()
def deploy(
    environment: Annotated[
        str, typer.Option(help="Environment to deploy to, e.g. dev, prod")
    ],
    service_name: Annotated[
        str, typer.Option(help="Name of the Airflow service to deploy")
    ],
    bucket: Annotated[
        str, typer.Option(help="S3 bucket name for deployment artifacts")
    ],
    kms_key_id: Annotated[
        str, typer.Option(help="KMS key ID to encrypt the deployment artifacts")
    ],
    git_commit: Annotated[str, typer.Option(help="Git commit hash for the deployment")],
    requirements_file: Annotated[
        str, typer.Option(help="Path to the requirements file")
    ] = "requirements.txt",
    dags_folder: Annotated[str, typer.Option(help="Path to the DAGs folder")] = "dags",
    bucket_prefix: Annotated[
        str, typer.Option(help="S3 bucket prefix for the deployment artifacts")
    ] = "",
    rollback: Annotated[
        bool, typer.Option(help="Use this to rollback to a previous deployment.")
    ] = False,
):
    bucket_prefix = (
        f"{bucket_prefix.strip('/')}/{git_commit}" if bucket_prefix else git_commit
    )
    requirements_key = f"{bucket_prefix}/requirements.txt"
    dags_key_prefix = f"{bucket_prefix}/dags"

    if not rollback:
        print(f"Deploying Airflow configuration to S3 bucket {bucket}")
        s3_connector = S3Connector()
        s3_connector.upload_file(
            bucket=bucket,
            filename=requirements_file,
            key=requirements_key,
            kms_key_id=kms_key_id,
        )
        s3_connector.upload_dir(
            bucket=bucket,
            folder=dags_folder,
            key_prefix=dags_key_prefix,
            kms_key_id=kms_key_id,
        )

    airflow_environment_name = f"{environment}-{service_name}"

    print(f"Updating airflow environment {airflow_environment_name}")
    mwaa_connector = MWAAConnector()
    mwaa_connector.update_environment(
        environment_name=airflow_environment_name,
        bucket=bucket,
        dag_s3_path=dags_key_prefix,
        requirements_s3_path=requirements_key,
    )

    print(
        f"Deployment to airflow environment {airflow_environment_name} "
        f"completed successfully."
    )