import os
import shlex
from kubernetes.client import models as k8s
from airflow.providers.docker.operators.docker import DockerOperator
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator

def get_container_operator(
    task_id: str,
    image: str,
    environment: dict,
    command: str,
    **kwargs
):
    """
    Returns a DockerOperator for local dev and a KubernetesPodOperator for prod.
    Controlled via the AIRFLOW_ENV environment variable.
    """
    airflow_env = os.getenv("AIRFLOW_ENV", "local")

    if airflow_env == "dev":
        return DockerOperator(
            task_id=task_id,
            image=image,
            api_version="auto",
            auto_remove="success",
            docker_url="unix://var/run/docker.sock",
            network_mode="airflow-docker_default",
            service_account_name="airflow-worker",
            mount_tmp_dir=False,
            environment=environment,
            command=command,
            **kwargs
        )
    else:
        # KPO prefers commands as lists. shlex safely splits the string.
        command_list = shlex.split(command) if isinstance(command, str) else command
        
        return KubernetesPodOperator(
            task_id=task_id,
            image=image,
            env_vars=environment,
            cmds=command_list,
            is_delete_operator_pod=True,
            in_cluster=True,
            get_logs=True,
            namespace="airflow",
            image_pull_secrets=[k8s.V1LocalObjectReference("cstg-pullsecret")],
            **kwargs
        )