"""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, State

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


@task(cache_policy=None)
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.",
)
@with_flow_timeout_param()
def transform_and_catalog(  # noqa: PLR0913, PLR0915
    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,
    additional_recipients: list[str] | None = None,
    recipients: list[str] | None = None,
    smtp_config: SmtpConfig | None = None,
    enable_notifications: bool = True,
    gh_action_actor: str | None = None,
) -> State | None:
    """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
        recipients (list[str] | None, optional): Primary recipient list. If provided,
            it takes precedence over the extracted owners email addresses from dbt
            metadata. Defaults to None.
        additional_recipients (list[str] | None, optional): Extra email addresses
            to be appended to the final recipient list regardless of other settings.
            Defaults to None.
        enable_notifications (bool): Whether to send email notifications for failed
            dbt tests. If True, notifications are sent to Technical Owners defined
            in dbt's meta config, with `notification_recipients` used as
            a fallback when no owner is found.
            If False, no notifications are sent regardless of
            `notification_recipients`. Defaults to True.
        gh_action_actor (str, optional): GitHub Actions actor that triggered the
            workflow. Defaults to None.

    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()
    if gh_action_actor:
        logger.info(f"Triggered by GitHub Actions actor: {gh_action_actor}")

    # Clone the dbt project.
    dbt_repo_url = dbt_repo_url or get_credentials(dbt_repo_url_secret)
    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.submit(
        project_path=dbt_project_path_full,
        command="deps",
    )
    pull_dbt_deps.result()

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

    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

            build = build_task.submit(
                project_path=dbt_project_path_full,
                command=f"build {build_select_safe} {dbt_target_option}",
                raise_on_failure=raise_on_failure,
                return_all=True,
            )
            build.result()
        else:
            run_task = dbt_task.with_options(name="dbt_run")
            run = run_task.submit(
                project_path=dbt_project_path_full,
                command=f"run {run_select_safe} {dbt_target_option}",
            )
            run.result()

            test_task = dbt_task.with_options(name="dbt_test")
            test = test_task.submit(
                project_path=dbt_project_path_full,
                command=f"test {test_select_safe} {dbt_target_option}",
                raise_on_failure=False,
            )
            test.result()
    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.submit(
            project_path=dbt_project_path_full,
            command="docs generate",
        )
        docs.result()

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

    upload_metadata = luma_ingest_task.submit(
        metadata_kind=metadata_kind,
        metadata_dir_path=dbt_target_dir_path,
        luma_url=luma_url,
        follow=luma_follow,
    )
    upload_metadata.result()

    file_name = "run_results.json"
    run_results_file_path = str(dbt_target_dir_path / file_name)  # type: ignore
    if run_results_storage_path:
        # Set the file path to include date info.
        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.

        s3_upload_file(
            from_path=run_results_file_path,
            to_path=run_results_storage_path,
            config_key=run_results_storage_config_key,
            credentials_secret=run_results_storage_credentials_secret,
        )
    if enable_notifications:
        dbt_test_failure_notifier(
            results_file_path=run_results_file_path,
            manifest_file_path=str(dbt_target_dir_path / "manifest.json"),
            recipients=recipients,
            additional_recipients=additional_recipients,
            smtp_config=smtp_config,
        )

    remove_dbt_repo_dir(
        dbt_repo_dir_name=dbt_repo_name,
    )

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

    return None
