"""
Shared Airflow task utilities for DockerOperator DAGs.

Place this file at dags/utils/common.py in the Airflow project.
"""

import string

from airflow.sdk import Variable
from airflow.decorators import task


@task()
def get_dlt_destination() -> str:
    """
    Read the DLT destination from Airflow Variable.
    """
    return Variable.get("dlt_destination")


@task(show_return_value_in_logs=False)
def get_dlt_secrets_toml_base64() -> str:
    """
    Read the DLT secrets TOML content from Airflow Variable and encode it in base64.
    """
    import base64
    import tomllib

    dlt_secrets_toml = Variable.get("dlt_secrets_toml")

    # Validate TOML content before encoding
    try:
        tomllib.loads(dlt_secrets_toml)
    except Exception as e:
        raise ValueError(f"Invalid TOML content in dlt_secrets_toml variable: {str(e)}")

    dlt_secrets_toml_base64 = base64.b64encode(dlt_secrets_toml.encode("utf-8")).decode(
        "utf-8"
    )
    return dlt_secrets_toml_base64


def sanitize_task_id(task_id: str) -> str:
    """Sanitize the task ID to ensure it is a valid Airflow task ID."""
    return (
        task_id.replace("-", "_").replace(" ", "_").replace(":", "_").replace("+", "_")
    )