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()

[((2696, 2750), 'metaflow.resources', 'resources', ([], {'local_storage': '"""242"""', 'cpu': '"""0.6"""', 'memory': '"""1G"""'}), "(local_storage='242', cpu='0.6', memory='1G')\n", (2705, 2750), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((2787, 2857), 'metaflow.environment', 'environment', ([], {'vars': "{'MY_ENV': 'value'}", 'kubernetes_vars': 'kubernetes_vars'}), "(vars={'MY_ENV': 'value'}, kubernetes_vars=kubernetes_vars)\n", (2798, 2857), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((4251, 4288), 'metaflow.environment', 'environment', ([], {'vars': "{'MY_ENV': 'value'}"}), "(vars={'MY_ENV': 'value'})\n", (4262, 4288), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((4319, 4342), 'metaflow.resources', 'resources', ([], {'volume': '"""11G"""'}), "(volume='11G')\n", (4328, 4342), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((4660, 4683), 'metaflow.resources', 'resources', ([], {'volume': '"""12G"""'}), "(volume='12G')\n", (4669, 4683), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((5049, 5101), 'metaflow.resources', 'resources', ([], {'volume': '"""13G"""', 'volume_mode': '"""ReadWriteMany"""'}), "(volume='13G', volume_mode='ReadWriteMany')\n", (5058, 5101), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), ((4488, 4560), 'subprocess.check_output', 'subprocess.check_output', (['"""df -h | grep /opt/metaflow_volume"""'], {'shell': '(True)'}), "('df -h | grep /opt/metaflow_volume', shell=True)\n", (4511, 4560), False, 'import subprocess\n'), ((4744, 4816), 'subprocess.check_output', 'subprocess.check_output', (['"""df -h | grep /opt/metaflow_volume"""'], {'shell': '(True)'}), "('df -h | grep /opt/metaflow_volume', shell=True)\n", (4767, 4816), False, 'import subprocess\n'), ((5171, 5243), 'subprocess.check_output', 'subprocess.check_output', (['"""df -h | grep /opt/metaflow_volume"""'], {'shell': '(True)'}), "('df -h | grep /opt/metaflow_volume', shell=True)\n", (5194, 5243), False, 'import subprocess\n'), ((3045, 3069), 'os.environ.get', 'os.environ.get', (['"""MY_ENV"""'], {}), "('MY_ENV')\n", (3059, 3069), False, 'import os\n'), ((3147, 3176), 'os.environ.get', 'os.environ.get', (['"""MY_POD_NAME"""'], {}), "('MY_POD_NAME')\n", (3161, 3176), False, 'import os\n'), ((3192, 3213), 'os.environ.get', 'os.environ.get', (['"""CPU"""'], {}), "('CPU')\n", (3206, 3213), False, 'import os\n'), ((3238, 3265), 'os.environ.get', 'os.environ.get', (['"""CPU_LIMIT"""'], {}), "('CPU_LIMIT')\n", (3252, 3265), False, 'import os\n'), ((3290, 3321), 'os.environ.get', 'os.environ.get', (['"""LOCAL_STORAGE"""'], {}), "('LOCAL_STORAGE')\n", (3304, 3321), False, 'import os\n'), ((3352, 3389), 'os.environ.get', 'os.environ.get', (['"""LOCAL_STORAGE_LIMIT"""'], {}), "('LOCAL_STORAGE_LIMIT')\n", (3366, 3389), False, 'import os\n'), ((3420, 3444), 'os.environ.get', 'os.environ.get', (['"""MEMORY"""'], {}), "('MEMORY')\n", (3434, 3444), False, 'import os\n'), ((3476, 3506), 'os.environ.get', 'os.environ.get', (['"""MEMORY_LIMIT"""'], {}), "('MEMORY_LIMIT')\n", (3490, 3506), False, 'import os\n'), ((3539, 3564), 'os.environ.get', 'os.environ.get', (['"""MF_NAME"""'], {}), "('MF_NAME')\n", (3553, 3564), False, 'import os\n'), ((3601, 3626), 'os.environ.get', 'os.environ.get', (['"""MF_STEP"""'], {}), "('MF_STEP')\n", (3615, 3626), False, 'import os\n'), ((3663, 3690), 'os.environ.get', 'os.environ.get', (['"""MF_RUN_ID"""'], {}), "('MF_RUN_ID')\n", (3677, 3690), False, 'import os\n'), ((3724, 3755), 'os.environ.get', 'os.environ.get', (['"""MF_EXPERIMENT"""'], {}), "('MF_EXPERIMENT')\n", (3738, 3755), False, 'import os\n'), ((3790, 3828), 'os.environ.get', 'os.environ.get', (['"""MF_TAG_METAFLOW_TEST"""'], {}), "('MF_TAG_METAFLOW_TEST')\n", (3804, 3828), False, 'import os\n'), ((3854, 3886), 'os.environ.get', 'os.environ.get', (['"""MF_TAG_TEST_T1"""'], {}), "('MF_TAG_TEST_T1')\n", (3868, 3886), False, 'import os\n'), ((3913, 3945), 'os.environ.get', 'os.environ.get', (['"""KF_POD_DEFAULT"""'], {}), "('KF_POD_DEFAULT')\n", (3927, 3945), False, 'import os\n'), ((3972, 4002), 'os.environ.get', 'os.environ.get', (['"""AI_FLOW_NAME"""'], {}), "('AI_FLOW_NAME')\n", (3986, 4002), False, 'import os\n'), ((4039, 4069), 'os.environ.get', 'os.environ.get', (['"""AI_STEP_NAME"""'], {}), "('AI_STEP_NAME')\n", (4053, 4069), False, 'import os\n'), ((4106, 4142), 'os.environ.get', 'os.environ.get', (['"""AI_EXPERIMENT_NAME"""'], {}), "('AI_EXPERIMENT_NAME')\n", (4120, 4142), False, 'import os\n'), ((4434, 4458), 'os.environ.get', 'os.environ.get', (['"""MY_ENV"""'], {}), "('MY_ENV')\n", (4448, 4458), False, 'import os\n'), ((5596, 5621), 'os.path.exists', 'os.path.exists', (['file_path'], {}), '(file_path)\n', (5610, 5621), False, 'import os\n'), ((5639, 5652), 'time.sleep', 'time.sleep', (['(1)'], {}), '(1)\n', (5649, 5652), False, 'import time\n'), ((1386, 1435), 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', ([], {'field_path': '"""metadata.name"""'}), "(field_path='metadata.name')\n", (1407, 1435), False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), ((1938, 2011), 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', ([], {'field_path': 'f"""metadata.annotations[\'{annotation}\']"""'}), '(field_path=f"metadata.annotations[\'{annotation}\']")\n', (1959, 2011), False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), ((2525, 2588), 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', ([], {'field_path': 'f"""metadata.labels[\'{label}\']"""'}), '(field_path=f"metadata.labels[\'{label}\']")\n', (2546, 2588), False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), ((694, 808), 'kubernetes.client.V1ResourceFieldSelector', 'V1ResourceFieldSelector', ([], {'container_name': '"""main"""', 'resource': 'resource', 'divisor': "('1m' if 'cpu' in resource else '1')"}), "(container_name='main', resource=resource, divisor=\n    '1m' if 'cpu' in resource else '1')\n", (717, 808), False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n')]