"""Build specified dbt model(s) and upload the generated metadata to Luma."""

from datetime import datetime, timezone
from pathlib import Path
import re
import shutil
from typing import Literal

from prefect import flow, task
from prefect.logging import get_run_logger
from prefect.states import Failed

from viadot.orchestration.prefect.tasks import (
    clone_repo,
    dbt_task,
    luma_ingest_task,
    s3_upload_file,
)
from viadot.orchestration.prefect.utils import get_credentials


@task
def remove_dbt_repo_dir(dbt_repo_dir_name: str) -> None:
    """Remove the repo directory.

    Args:
        dbt_repo_dir_name (str): The name of the dbt repo directory.
    """
    shutil.rmtree(dbt_repo_dir_name, ignore_errors=True)


@flow(
    name="Transform and Catalog",
    description="Build specified dbt model(s) and upload generated metadata to Luma.",
    timeout_seconds=2 * 60 * 60,
)
def transform_and_catalog(  # noqa: PLR0912, PLR0913, PLR0915, C901
    dbt_repo_url: str | None = None,
    dbt_repo_url_secret: str | None = None,
    dbt_project_path: str = "dbt",
    dbt_repo_branch: str | None = None,
    dbt_repo_token_secret: str | None = None,
    dbt_selects: dict[str, str] | None = None,
    dbt_target: str | None = None,
    dbt_target_dir_path: str | None = None,
    luma_url: str = "http://localhost:8000",
    luma_follow: bool = False,
    metadata_kind: Literal["model", "model_run"] = "model_run",
    run_results_storage_path: str | None = None,
    run_results_storage_config_key: str | None = None,
    run_results_storage_credentials_secret: str | None = None,
    fail_flow_only_on_build_failure: bool = False,
) -> list[str]:
    """Build specified dbt model(s) and upload the generated metadata to Luma.

    Supports ingesting both model and model run metadata (controlled by the
    `metadata_kind` parameter).

    Note that metadata is still ingested even if the preceding `dbt test` task fails.
    This is done in order to capture test failure metadata in the data catalog.

    Args:
        dbt_repo_url (str, optional): The URL for cloning the dbt repo with relevant
            dbt project. Defaults to None.
        dbt_repo_url_secret (str, optional): Alternatively to above, the secret
            containing `dbt_repo_url`. Defaults to None.
        dbt_project_path (str): Path to the dbt project directory, relative to the
            dbt repository's root. For example, "dbt/my_dbt_project". Defaults to "dbt".
        dbt_repo_branch (str, optional): The branch of the dbt repo to use. Defaults to
            None (default repo branch).
        dbt_repo_token_secret (str, optional): The secret containing the personal access
            token used to clone the dbt repository, in case it's private. Not required
            if token is already included in `dbt_repo_url` (which is NOT recommended).
            Defaults to None.
        dbt_selects (dict, optional): Valid
            [dbt node selection](https://docs.getdbt.com/reference/node-selection/syntax)
            expressions. Valid keys are `run`, `test`,`build`, and `source_freshness`.
                The test select expression is taken from run's, as long as run select is
                provided. Defaults to None.
        dbt_target (str): The dbt target to use. If not specified, the default dbt
            target (as specified in `profiles.yaml`) will be used. Defaults to None.
        dbt_target_dir_path (str): The path to your dbt project's target
            directory, which contains dbt artifact JSON files, relative
            to dbt project's root directory. By default,
            `<repo_name>/<dbt_project_path>/target`, since "target" is the default
            name of the directory generated by dbt.
        luma_url (str, optional): The URL of the Luma instance to ingest into.
            Defaults to "http://localhost:8000".
        luma_follow (bool, optional): Whether to follow the ingestion process until it's
            completed (by default, ingestion request is sent without awaiting for the
            response). By default, `False`.
        metadata_kind (Literal["model", "model_run"], optional): The kind of metadata
            to ingest. Defaults to "model_run".
        run_results_storage_path (str, optional): The directory to upload the
            `run_results.json` file to. Note that a timestamp will be appended to the
            end of the file. Currently, only S3 is supported. Defaults to None.
        run_results_storage_config_key (str, optional): The key in the viadot config
            holding AWS credentials. Defaults to None.
        run_results_storage_credentials_secret (str, optional): The name of the secret
            block in Prefect holding AWS credentials. Defaults to None.
        fail_flow_only_on_build_failure (bool): Determines the flow's failure behavior
            based on dbt build outcomes.
            When False (default):
                - The flow will fail on any dbt build failure (including test failures)
            When True:
                - The flow will only fail if model building fails
                - Test failures alone won't cause the flow failure

    Returns:
        list[str]: Lines from stdout of the `upload_metadata` task as a list.

    Examples:
        # Build staging models.

        ```python
        import os
        from prefect_viadot.flows import transform_and_catalog

        my_dbt_repo_url = "https://github.com/dbt-labs/jaffle_shop"
        my_luma_url = "http://localhost:8000"

        transform_and_catalog(
            dbt_repo_url=my_dbt_repo_url
            dbt_selects={"run": "staging"}
            luma_url=my_luma_url,
            run_results_storage_path="s3://my-bucket/dbt/run_results",
            run_results_storage_credentials_secret="my-aws-credentials-block",
        )
        ```

        Some common `dbt_selects` patterns:
        - runs a specific model and all its downstream dependencies:
            `dbt_select={"run": "my_model+"}`
        - runs all models in a directory:
            `dbt_select={"run: "models/staging"}`
        - runs a specific model in a folder:
            `dbt_select={"run": "marts.domain.some_model"}`
        - runs tests for a specific model:
            `dbt_select={"test": "my_model"}`
        - build a specific model:
            `dbt_select={"build": "my_model"}`
        - build all models in a folder:
            `dbt_select={"build": "models.intermediate"}`
    """
    logger = get_run_logger()
    # Clone the dbt project.
    dbt_repo_url = dbt_repo_url or get_credentials(dbt_repo_url_secret)
    clone = clone_repo(
        url=dbt_repo_url,
        checkout_branch=dbt_repo_branch,
        token_secret=dbt_repo_token_secret,
    )

    # Prepare the environment.
    dbt_repo_name = dbt_repo_url.split("/")[-1].replace(".git", "")
    dbt_project_path_full = Path(dbt_repo_name) / dbt_project_path
    dbt_pull_deps_task = dbt_task.with_options(
        name="dbt_deps", retries=3, retry_delay_seconds=60
    )
    pull_dbt_deps = dbt_pull_deps_task(
        project_path=dbt_project_path_full,
        command="deps",
        wait_for=[clone],
    )

    # Run dbt commands.
    dbt_target_option = f"-t {dbt_target}" if dbt_target is not None else ""

    task_failed = False

    if metadata_kind == "model_run":
        # Produce `run-results.json` artifact for Luma ingestion.
        if dbt_selects:
            build_select = dbt_selects.get("build")
            run_select = dbt_selects.get("run")
            test_select = dbt_selects.get("test", run_select)

            build_select_safe = f"-s {build_select}" if build_select is not None else ""
            run_select_safe = f"-s {run_select}" if run_select is not None else ""
            test_select_safe = f"-s {test_select}" if test_select is not None else ""
        else:
            run_select_safe = ""
            test_select_safe = ""
        if build_select:
            # If build task is used, run and test tasks are not needed.
            # Build task executes run and tests commands internally.
            build_task = dbt_task.with_options(name="dbt_build")
            raise_on_failure = not fail_flow_only_on_build_failure
            try:
                build = build_task(
                    project_path=dbt_project_path_full,
                    command=f"build {build_select_safe} {dbt_target_option}",
                    wait_for=[pull_dbt_deps],
                    raise_on_failure=raise_on_failure,
                    return_all=True,
                )
            except Exception:
                msg = "Build task failed."
                logger.exception(msg)
                build = msg
                task_failed = True

            upload_metadata_upstream_task = build
        else:
            run_task = dbt_task.with_options(name="dbt_run")
            run = run_task(
                project_path=dbt_project_path_full,
                command=f"run {run_select_safe} {dbt_target_option}",
                wait_for=[pull_dbt_deps],
            )

            test_task = dbt_task.with_options(name="dbt_test")
            test = test_task(
                project_path=dbt_project_path_full,
                command=f"test {test_select_safe} {dbt_target_option}",
                raise_on_failure=False,
                wait_for=[run],
            )
            upload_metadata_upstream_task = test

    else:
        # Produce `catalog.json` and `manifest.json` artifacts for Luma ingestion.
        docs_generate_task = dbt_task.with_options(name="dbt_docs_generate")
        docs = docs_generate_task(
            project_path=dbt_project_path_full,
            command="docs generate",
            wait_for=[pull_dbt_deps],
        )
        upload_metadata_upstream_task = docs

    # Upload metadata to Luma.
    if dbt_target_dir_path is None:
        dbt_target_dir_path = dbt_project_path_full / "target"

    try:
        upload_metadata = luma_ingest_task(
            metadata_kind=metadata_kind,
            metadata_dir_path=dbt_target_dir_path,
            luma_url=luma_url,
            follow=luma_follow,
            wait_for=[upload_metadata_upstream_task],
        )
    except Exception:
        msg = "Luma ingest task failed."
        logger.exception(msg)
        upload_metadata = msg
        task_failed = True

    if run_results_storage_path:
        # Set the file path to include date info.
        file_name = "run_results.json"
        now = datetime.now(timezone.utc)
        run_results_storage_path = run_results_storage_path.rstrip("/") + "/"

        # Add partitioning.
        date_str = now.strftime("%Y%m%d")
        run_results_storage_path += date_str + "/"

        # Add timestamp suffix, eg. run_results_1737556947.934292.json.
        timestamp = now.timestamp()
        run_results_storage_path += (
            Path(file_name).stem + "_" + str(timestamp) + ".json"
        )
        logger.info(f"Uploading run results to {run_results_storage_path}")
        # Upload the file to s3.
        try:
            dump_test_results_to_s3 = s3_upload_file(
                from_path=str(dbt_target_dir_path / file_name),
                to_path=run_results_storage_path,
                wait_for=[upload_metadata_upstream_task],
                config_key=run_results_storage_config_key,
                credentials_secret=run_results_storage_credentials_secret,
            )
        except Exception:
            msg = "S3 upload file task failed."
            logger.exception(msg)
            dump_test_results_to_s3 = msg
            task_failed = True

    # Cleanup.
    wait_for = (
        [upload_metadata, dump_test_results_to_s3]
        if run_results_storage_path
        else [upload_metadata]
    )
    remove_dbt_repo_dir(dbt_repo_name, wait_for=wait_for)

    if task_failed:
        return Failed()

    if fail_flow_only_on_build_failure and build_select:
        model_error_pattern = re.compile(r"ERROR creating", re.IGNORECASE)
        if any(model_error_pattern.search(line) for line in build):
            return Failed(message="One or more models failed to build.")

    return remove_dbt_repo_dir
