"""
Shared Airflow utility for running dbt commands via DockerOperator.

Place this file at dags/utils/dbt.py in the Airflow project.
Requires the following environment variables (or Airflow .env):
  - AIRFLOW_DBT_ADAPTER_TYPE (default: postgres)
  - AIRFLOW_DBT_ADAPTER_VERSION (default: 1.9)
  - AIRFLOW_DBT_PROJECT_PATH - path to the dbt project on the host
  - AIRFLOW_DBT_PROFILES_PATH - path to profiles.yml on the host
  - AIRFLOW_DBT_PRIVATE_KEY_PATH - path to private key on the host (for Snowflake)
"""

import os

from airflow.providers.docker.operators.docker import DockerOperator
from docker.types import Mount

from utils.common import sanitize_task_id


DBT_ADAPTER_TYPE = os.getenv("AIRFLOW_DBT_ADAPTER_TYPE", "postgres")
DBT_ADAPTER_VERSION = os.getenv("AIRFLOW_DBT_ADAPTER_VERSION", "1.9")
DBT_DOCKER_IMAGE = (
    f"ghcr.io/dbt-labs/dbt-{DBT_ADAPTER_TYPE}:{DBT_ADAPTER_VERSION}.latest"
)

DBT_PROJECT_PATH = os.getenv(
    "AIRFLOW_DBT_PROJECT_PATH", "/home/gemma/repos/data-transformations"
)
DBT_PROFILES_PATH = os.getenv(
    "AIRFLOW_DBT_PROFILES_PATH",
    "/home/gemma/.dbt/profiles.yml",
)
DBT_PRIVATE_KEY_PATH = os.getenv(
    "AIRFLOW_DBT_PRIVATE_KEY_PATH", "/home/gemma/.dbt/dbt_snowflake.p8"
)


def create_dbt_task(command: str, task_id: str = None):
    """
    Create a dbt task that runs a dbt command using a Docker container.

    Args:
        command: The dbt command to run (e.g., "run", "test", "build").
        task_id: Optional custom task ID. Defaults to "dbt_<command>".

    Returns:
        A DockerOperator task that runs the dbt command.
    """
    if not task_id:
        task_id = sanitize_task_id(f"dbt_{command}")
    return DockerOperator(
        task_id=task_id,
        image=DBT_DOCKER_IMAGE,
        container_name=task_id,
        api_version="auto",
        auto_remove="force",
        docker_url="tcp://docker-proxy:2375",
        mount_tmp_dir=False,
        mounts=[
            Mount(
                target="/root/.dbt/profiles.yml",
                source=DBT_PROFILES_PATH,
                type="bind",
            ),
            Mount(
                target="/usr/app",
                source=DBT_PROJECT_PATH,
                type="bind",
            ),
            Mount(
                target="/root/.dbt/private_key.p8",
                source=DBT_PRIVATE_KEY_PATH,
                type="bind",
            ),
        ],
        command=command,
    )