import pendulum
import base64
from airflow.sdk import dag, task
from airflow.providers.cncf.kubernetes.operators.pod import (
    KubernetesPodOperator,
)
from stackit_workflows.kubernetes import K8S, V1, client
from airflow.settings import json
from kubernetes.client import models as k8s
# It is good practice but not required to specify the image to use.
default_kwargs = {
    "image": "schwarzit-xx-sit-dp-customer-artifactory-docker-local.jfrog.io/stackit-spark:spark3.5.3-0.1.2"
}


@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["demo","stackit-demo"],
    dag_id="stackit_11_custom_secret_podoperator",
)
def simple_spark():

    k = KubernetesPodOperator(
        name="hello-dry-run",
        image="schwarzit-xx-sit-ragstudio-docker-local.jfrog.io/airflowrunner:latest",
        image_pull_secrets=[K8S.V1LocalObjectReference("internal-dockerpull-secret")],
        cmds=["bash", "-cx"],
        arguments=["echo", "10"],
        labels={"foo": "bar"},
        task_id="dry_run_demo",
        log_events_on_failure=True,
        do_xcom_push=True,
        #Container resources requests and limits. Mandatory to set
        container_resources=K8S.V1ResourceRequirements(
            requests={"cpu": "100m", "memory": "50Mi"}, #optional: "ephemeral-storage": "1Gi"
            limits={"cpu": "200m", "memory": "100Mi"},  #optional "ephemeral-storage": "2Gi"
        ),
        # All pods launched must run as non-root users. Otherwise they won't start.
        security_context=K8S.V1PodSecurityContext(run_as_user=100),
    )

    return k

simple_spark()