import os
import pprint
import subprocess
import time
from typing import Dict, List

from kubernetes.client import (
    V1EnvVar,
    V1EnvVarSource,
    V1ObjectFieldSelector,
    V1ResourceFieldSelector,
)

from metaflow import FlowSpec, step, environment, resources, current


def get_env_vars(env_resources: Dict[str, str]) -> List[V1EnvVar]:
    res = []
    for name, resource in env_resources.items():
        res.append(
            V1EnvVar(
                # this is used by some functions of operator-sdk
                # it uses this environment variable to get the pods
                name=name,
                value_from=V1EnvVarSource(
                    resource_field_ref=V1ResourceFieldSelector(
                        container_name="main",
                        resource=resource,
                        divisor="1m" if "cpu" in resource else "1",
                    )
                ),
            )
        )
    return res


kubernetes_vars = get_env_vars(
    {
        "LOCAL_STORAGE": "requests.ephemeral-storage",
        "LOCAL_STORAGE_LIMIT": "limits.ephemeral-storage",
        "CPU": "requests.cpu",
        "CPU_LIMIT": "limits.cpu",
        "MEMORY": "requests.memory",
        "MEMORY_LIMIT": "limits.memory",
    }
)
kubernetes_vars.append(
    V1EnvVar(
        name="MY_POD_NAME",
        value_from=V1EnvVarSource(
            field_ref=V1ObjectFieldSelector(field_path="metadata.name")
        ),
    )
)

annotations = {
    "metaflow.org/flow_name": "MF_NAME",
    "metaflow.org/step": "MF_STEP",
    "metaflow.org/run_id": "MF_RUN_ID",
    "metaflow.org/experiment": "MF_EXPERIMENT",
    "metaflow.org/tag_metaflow_test": "MF_TAG_METAFLOW_TEST",
    "metaflow.org/tag_test_t1": "MF_TAG_TEST_T1",
}
for annotation, env_name in annotations.items():
    kubernetes_vars.append(
        V1EnvVar(
            name=env_name,
            value_from=V1EnvVarSource(
                field_ref=V1ObjectFieldSelector(
                    field_path=f"metadata.annotations['{annotation}']"
                )
            ),
        )
    )

labels = {
    "aip.zillowgroup.net/kfp-pod-default": "KF_POD_DEFAULT",
    "tags.ledger.zgtools.net/ai-flow-name": "AI_FLOW_NAME",
    "tags.ledger.zgtools.net/ai-step-name": "AI_STEP_NAME",
    "tags.ledger.zgtools.net/ai-experiment-name": "AI_EXPERIMENT_NAME",
}
for label, env_name in labels.items():
    kubernetes_vars.append(
        V1EnvVar(
            name=env_name,
            value_from=V1EnvVarSource(
                field_ref=V1ObjectFieldSelector(
                    field_path=f"metadata.labels['{label}']"
                )
            ),
        )
    )


class ResourcesFlow(FlowSpec):
    @resources(
        local_storage="242",
        cpu="0.6",
        memory="1G",
    )
    @environment(  # pylint: disable=E1102
        vars={"MY_ENV": "value"}, kubernetes_vars=kubernetes_vars
    )
    @step
    def start(self):
        pprint.pprint(dict(os.environ))
        print("=====")

        # test simple environment var
        assert os.environ.get("MY_ENV") == "value"

        # test kubernetes_vars
        assert "resourcesflow" in os.environ.get("MY_POD_NAME")
        assert os.environ.get("CPU") == "600"
        assert os.environ.get("CPU_LIMIT") == "600"
        assert os.environ.get("LOCAL_STORAGE") == "242000000"
        assert os.environ.get("LOCAL_STORAGE_LIMIT") == "242000000"
        assert os.environ.get("MEMORY") == "1000000000"
        assert os.environ.get("MEMORY_LIMIT") == "1000000000"

        assert os.environ.get("MF_NAME") == current.flow_name
        assert os.environ.get("MF_STEP") == current.step_name
        assert os.environ.get("MF_RUN_ID") == current.run_id
        assert os.environ.get("MF_EXPERIMENT") == "metaflow_test"
        assert os.environ.get("MF_TAG_METAFLOW_TEST") == "true"
        assert os.environ.get("MF_TAG_TEST_T1") == "true"

        assert os.environ.get("KF_POD_DEFAULT") == "true"

        assert os.environ.get("AI_FLOW_NAME") == current.flow_name
        assert os.environ.get("AI_STEP_NAME") == current.step_name
        assert os.environ.get("AI_EXPERIMENT_NAME") == "metaflow_test"

        self.items = [1, 2]
        self.next(self.foreach_step, foreach="items")

    @environment(vars={"MY_ENV": "value"})  # pylint: disable=E1102
    @resources(volume="11G")
    @step
    def foreach_step(self):
        # test simple environment var
        assert os.environ.get("MY_ENV") == "value"

        output = subprocess.check_output(
            "df -h | grep /opt/metaflow_volume", shell=True
        )
        assert "11G" in str(output)

        self.next(self.join_step)

    @resources(volume="12G")
    @step
    def join_step(self, inputs):
        output = subprocess.check_output(
            "df -h | grep /opt/metaflow_volume", shell=True
        )
        assert "12G" in str(output)
        self.next(self.split_step)

    @step
    def split_step(self):
        self.items = [1, 2]
        self.next(self.shared_volume_foreach_step, foreach="items")

    @resources(volume="13G", volume_mode="ReadWriteMany")
    @step
    def shared_volume_foreach_step(self):
        output = subprocess.check_output(
            "df -h | grep /opt/metaflow_volume", shell=True
        )
        assert "13G" in str(output)

        file_path = "/opt/metaflow_volume/test.txt"
        message = "hello world!"

        # validate the volume is shared across the foreach splits
        if self.input == 1:
            with open(file_path, "w") as f:
                f.write(message)
        else:
            while not os.path.exists(file_path):
                time.sleep(1)
                print(".")

            with open(file_path, "r") as f:
                read_lines = f.readlines()
                print("read_lines", read_lines)
                assert message == read_lines[0]

        self.next(self.shared_volume_join_step)

    @step
    def shared_volume_join_step(self, inputs):
        self.next(self.end)

    @step
    def end(self):
        print("All done.")


if __name__ == "__main__":
    ResourcesFlow()

[((96, 5, 100, 5), 'metaflow.resources', 'resources', (), '', False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((101, 5, 103, 5), 'metaflow.environment', 'environment', (), '', False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((137, 5, 137, 42), 'metaflow.environment', 'environment', (), '', False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((138, 5, 138, 28), 'metaflow.resources', 'resources', (), '', False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((151, 5, 151, 28), 'metaflow.resources', 'resources', (), '', False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((165, 5, 165, 57), 'metaflow.resources', 'resources', (), '', False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((144, 17, 146, 9), 'subprocess.check_output', 'subprocess.check_output', (), '', False, 'import subprocess\n'), ((154, 17, 156, 9), 'subprocess.check_output', 'subprocess.check_output', (), '', False, 'import subprocess\n'), ((168, 17, 170, 9), 'subprocess.check_output', 'subprocess.check_output', (), '', False, 'import subprocess\n'), ((110, 15, 110, 39), 'os.environ.get', 'os.environ.get', ({(110, 30, 110, 38): '"""MY_ENV"""'}, {}), "('MY_ENV')", False, 'import os\n'), ((113, 34, 113, 63), 'os.environ.get', 'os.environ.get', ({(113, 49, 113, 62): '"""MY_POD_NAME"""'}, {}), "('MY_POD_NAME')", False, 'import os\n'), ((114, 15, 114, 36), 'os.environ.get', 'os.environ.get', ({(114, 30, 114, 35): '"""CPU"""'}, {}), "('CPU')", False, 'import os\n'), ((115, 15, 115, 42), 'os.environ.get', 'os.environ.get', ({(115, 30, 115, 41): '"""CPU_LIMIT"""'}, {}), "('CPU_LIMIT')", False, 'import os\n'), ((116, 15, 116, 46), 'os.environ.get', 'os.environ.get', ({(116, 30, 116, 45): '"""LOCAL_STORAGE"""'}, {}), "('LOCAL_STORAGE')", False, 'import os\n'), ((117, 15, 117, 52), 'os.environ.get', 'os.environ.get', ({(117, 30, 117, 51): '"""LOCAL_STORAGE_LIMIT"""'}, {}), "('LOCAL_STORAGE_LIMIT')", False, 'import os\n'), ((118, 15, 118, 39), 'os.environ.get', 'os.environ.get', ({(118, 30, 118, 38): '"""MEMORY"""'}, {}), "('MEMORY')", False, 'import os\n'), ((119, 15, 119, 45), 'os.environ.get', 'os.environ.get', ({(119, 30, 119, 44): '"""MEMORY_LIMIT"""'}, {}), "('MEMORY_LIMIT')", False, 'import os\n'), ((121, 15, 121, 40), 'os.environ.get', 'os.environ.get', ({(121, 30, 121, 39): '"""MF_NAME"""'}, {}), "('MF_NAME')", False, 'import os\n'), ((122, 15, 122, 40), 'os.environ.get', 'os.environ.get', ({(122, 30, 122, 39): '"""MF_STEP"""'}, {}), "('MF_STEP')", False, 'import os\n'), ((123, 15, 123, 42), 'os.environ.get', 'os.environ.get', ({(123, 30, 123, 41): '"""MF_RUN_ID"""'}, {}), "('MF_RUN_ID')", False, 'import os\n'), ((124, 15, 124, 46), 'os.environ.get', 'os.environ.get', ({(124, 30, 124, 45): '"""MF_EXPERIMENT"""'}, {}), "('MF_EXPERIMENT')", False, 'import os\n'), ((125, 15, 125, 53), 'os.environ.get', 'os.environ.get', ({(125, 30, 125, 52): '"""MF_TAG_METAFLOW_TEST"""'}, {}), "('MF_TAG_METAFLOW_TEST')", False, 'import os\n'), ((126, 15, 126, 47), 'os.environ.get', 'os.environ.get', ({(126, 30, 126, 46): '"""MF_TAG_TEST_T1"""'}, {}), "('MF_TAG_TEST_T1')", False, 'import os\n'), ((128, 15, 128, 47), 'os.environ.get', 'os.environ.get', ({(128, 30, 128, 46): '"""KF_POD_DEFAULT"""'}, {}), "('KF_POD_DEFAULT')", False, 'import os\n'), ((130, 15, 130, 45), 'os.environ.get', 'os.environ.get', ({(130, 30, 130, 44): '"""AI_FLOW_NAME"""'}, {}), "('AI_FLOW_NAME')", False, 'import os\n'), ((131, 15, 131, 45), 'os.environ.get', 'os.environ.get', ({(131, 30, 131, 44): '"""AI_STEP_NAME"""'}, {}), "('AI_STEP_NAME')", False, 'import os\n'), ((132, 15, 132, 51), 'os.environ.get', 'os.environ.get', ({(132, 30, 132, 50): '"""AI_EXPERIMENT_NAME"""'}, {}), "('AI_EXPERIMENT_NAME')", False, 'import os\n'), ((142, 15, 142, 39), 'os.environ.get', 'os.environ.get', ({(142, 30, 142, 38): '"""MY_ENV"""'}, {}), "('MY_ENV')", False, 'import os\n'), ((181, 22, 181, 47), 'os.path.exists', 'os.path.exists', ({(181, 37, 181, 46): 'file_path'}, {}), '(file_path)', False, 'import os\n'), ((182, 16, 182, 29), 'time.sleep', 'time.sleep', ({(182, 27, 182, 28): '(1)'}, {}), '(1)', False, 'import time\n'), ((51, 22, 51, 71), 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', (), '', False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), ((69, 26, 71, 17), 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', (), '', False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), ((87, 26, 89, 17), 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', (), '', False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), ((26, 39, 30, 21), 'kubernetes.client.V1ResourceFieldSelector', 'V1ResourceFieldSelector', (), '', False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n')]