import logging
import os
import shutil
from pathlib import Path
from typing import Union

import boto3
import cloudpickle as pickle

from conductor.docker import CONTAINER_DIR, CONTAINER_FILES
from conductor.internal.utils import (
    exec_sync,
    get_dags_s3_path,
    get_project_files,
    load_conductor_instances,
    load_project,
)

SERIALIZED_DAG_TEMPLATE = Path(
    os.path.dirname(__file__), "templates", "serialized_dag.py.txt"
)


def deploy(project_root, selected_dag_name=None, remote=True) -> None:
    # Ensure build directory is a directory.
    build_dir = clean_build(project_root)

    # Set CDK_OUTDIR before projects are loaded.
    cdk_outdir = Path(build_dir, "cdk")
    os.environ["CDK_OUTDIR"] = str(cdk_outdir)

    # Load the project.
    project = load_project(project_root)
    conductor_instances = load_conductor_instances(project_root)

    # Compile DAGs
    dag_files = compile_dags(
        conductor_instances, build_dir, selected_dag_name=selected_dag_name
    )

    # Build Docker Image
    ecr_url = project.project_resources.ecr_url()
    build_image(ecr_url, build_dir)

    if remote:
        # Set the AWS_PROFILE environment variable for AWS commands.
        os.environ["AWS_PROFILE"] = project.project_resources.env.account_name

        # Deploy CDK resources.
        project.project_resources.synth()
        exec_sync(
            f"{project.project_resources.cdk_executable_path()} --app {cdk_outdir} deploy --require-approval never"
        )
        # Push Docker Image
        push_image(ecr_url)
        # Push the serialized DAGs to the target environment
        push_dags(project, dag_files)


def clean_build(project_root):
    build_dir = Path(project_root, "build")
    if build_dir.exists():
        if not build_dir.is_dir():
            raise Exception(f"{build_dir} exists but is not a directory")
    else:
        build_dir.mkdir()
    return build_dir


def compile_dags(conductor_instances, build_dir, selected_dag_name: str = None):
    # Refresh the DAG build directory.
    dags_output_dir = Path(build_dir, "serialized_dags")
    shutil.rmtree(dags_output_dir, ignore_errors=True)
    dags_output_dir.mkdir()

    dag_files = []

    for c in conductor_instances:
        if selected_dag_name is not None:
            if c.name == selected_dag_name:
                pass
            else:
                continue
        serialized_dag = pickle.dumps(c.dag)
        dag_file_contents = (
            open(SERIALIZED_DAG_TEMPLATE, "r")
            .read()
            .format(serialized_dag=serialized_dag)
        )
        dag_file = Path(dags_output_dir, f"{c.dag.dag_id}.py")
        dag_files.append(dag_file)
        with open(dag_file, "w+") as f:
            f.write(dag_file_contents)
    if selected_dag_name is not None and len(dag_files) == 0:
        raise ValueError(
            f"Selected DAG {selected_dag_name} was not found among the DAGs in this project."
        )
    return dag_files


def push_dags(project, dag_files) -> None:
    airflow_config = project.project_resources.env.airflow
    bucket_name, dag_s3_path = get_dags_s3_path(airflow_config)
    if bucket_name is None or dag_s3_path is None:
        logging.warning("No Airflow remote specified, skip pushing DAGs.")
        return
    s3 = boto3.resource("s3")
    for dag_file in dag_files:
        s3.meta.client.upload_file(
            str(dag_file), bucket_name, str(Path(dag_s3_path, dag_file.name))
        )


def build_image(ecr_url: str, build_dir: Union[str, Path]):
    # Clean if already exists
    docker_build_dir = Path(build_dir, "docker")
    shutil.rmtree(str(docker_build_dir), ignore_errors=True)
    docker_build_dir.mkdir()

    # Copy container entrypoint files
    for f in CONTAINER_FILES:
        target_path = Path(docker_build_dir, f)
        target_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(Path(CONTAINER_DIR, f), target_path)

    # Copy user files not ignored by Git to build dir.
    for f in get_project_files():
        target_path = Path(docker_build_dir, f)
        target_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(str(f), target_path)

    exec_sync(f"docker build {docker_build_dir} -t {ecr_url}")


def push_image(ecr_url: str):
    exec_sync(
        "aws ecr get-login-password --region us-west-2 | "
        f"docker login --username AWS --password-stdin {ecr_url}"
    )
    exec_sync(f"docker push {ecr_url}")
