import os
from prefect import task, flow, get_run_logger

from prefect.blocks.system import Secret
from prefect_dbt.cli import DbtCliProfile, DbtCoreOperation
from prefect_shell import ShellOperation


ENV = os.environ["ENV"].upper() if "ENV" in os.environ else "DEV"
DBT_ENV = ENV.lower()


@task(retries=10, retry_delay_seconds=60)
def checkout_dbt_repo(token: str):
    ShellOperation(
        commands=[
            "rm -rf edw-dbt && git clone --depth=1 --branch=main https://{github_access_token}@github.com/cloudfactory/edw-dbt".format(
                github_access_token=token
            )
        ]
    ).run()


@flow(name="dbt_flow")
def dbt_flow(dbt_build_arg: str = ""):
    logger = get_run_logger()
    <http://logger.info|logger.info>(f"flow parameter dbt_build_arg has value : {dbt_build_arg}")
    gh_access_token = Secret.load("github-access-token").get()

    <http://logger.info|logger.info>("Downloading dbt repo")
    checkout_dbt_repo(token=gh_access_token)

    dbt_cli_profile = DbtCliProfile.load("dbt-cli-profile")
    <http://logger.info|logger.info>(f"Running - dbt build {dbt_build_arg} --target {DBT_ENV}")
    with DbtCoreOperation(
        commands=["dbt debug", "dbt deps", "dbt build " + dbt_build_arg + " --target " + DBT_ENV],
        project_dir="edw-dbt",
        overwrite_profiles=True,
        dbt_cli_profile=dbt_cli_profile,
        stream_output=True,
    ) as dbt_operation:
        dbt_process = dbt_operation.trigger()
        dbt_process.wait_for_completion()
        result = dbt_process.fetch_result()

    return result


if __name__ == "__main__":
    print(dbt_flow("-s base_users_users"))