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, 'metaflow.resources', 'resources', ([], {'local_storage': '"""242"""', 'cpu': '"""0.6"""', 'memory': '"""1G"""'}), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), (101, 'metaflow.environment', 'environment', ([], {'vars': "{'MY_ENV': 'value'}", 'kubernetes_vars': 'kubernetes_vars'}), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), (137, 'metaflow.environment', 'environment', ([], {'vars': "{'MY_ENV': 'value'}"}), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), (138, 'metaflow.resources', 'resources', ([], {'volume': '"""11G"""'}), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), (151, 'metaflow.resources', 'resources', ([], {'volume': '"""12G"""'}), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), (165, 'metaflow.resources', 'resources', ([], {'volume': '"""13G"""', 'volume_mode': '"""ReadWriteMany"""'}), False, 'from metaflow import FlowSpec, step, environment, resources, current\n'), (144, 'subprocess.check_output', 'subprocess.check_output', (['"""df -h | grep /opt/metaflow_volume"""'], {'shell': '(True)'}), False, 'import subprocess\n'), (154, 'subprocess.check_output', 'subprocess.check_output', (['"""df -h | grep /opt/metaflow_volume"""'], {'shell': '(True)'}), False, 'import subprocess\n'), (168, 'subprocess.check_output', 'subprocess.check_output', (['"""df -h | grep /opt/metaflow_volume"""'], {'shell': '(True)'}), False, 'import subprocess\n'), (110, 'os.environ.get', 'os.environ.get', (['"""MY_ENV"""'], {}), False, 'import os\n'), (113, 'os.environ.get', 'os.environ.get', (['"""MY_POD_NAME"""'], {}), False, 'import os\n'), (114, 'os.environ.get', 'os.environ.get', (['"""CPU"""'], {}), False, 'import os\n'), (115, 'os.environ.get', 'os.environ.get', (['"""CPU_LIMIT"""'], {}), False, 'import os\n'), (116, 'os.environ.get', 'os.environ.get', (['"""LOCAL_STORAGE"""'], {}), False, 'import os\n'), (117, 'os.environ.get', 'os.environ.get', (['"""LOCAL_STORAGE_LIMIT"""'], {}), False, 'import os\n'), (118, 'os.environ.get', 'os.environ.get', (['"""MEMORY"""'], {}), False, 'import os\n'), (119, 'os.environ.get', 'os.environ.get', (['"""MEMORY_LIMIT"""'], {}), False, 'import os\n'), (121, 'os.environ.get', 'os.environ.get', (['"""MF_NAME"""'], {}), False, 'import os\n'), (122, 'os.environ.get', 'os.environ.get', (['"""MF_STEP"""'], {}), False, 'import os\n'), (123, 'os.environ.get', 'os.environ.get', (['"""MF_RUN_ID"""'], {}), False, 'import os\n'), (124, 'os.environ.get', 'os.environ.get', (['"""MF_EXPERIMENT"""'], {}), False, 'import os\n'), (125, 'os.environ.get', 'os.environ.get', (['"""MF_TAG_METAFLOW_TEST"""'], {}), False, 'import os\n'), (126, 'os.environ.get', 'os.environ.get', (['"""MF_TAG_TEST_T1"""'], {}), False, 'import os\n'), (128, 'os.environ.get', 'os.environ.get', (['"""KF_POD_DEFAULT"""'], {}), False, 'import os\n'), (130, 'os.environ.get', 'os.environ.get', (['"""AI_FLOW_NAME"""'], {}), False, 'import os\n'), (131, 'os.environ.get', 'os.environ.get', (['"""AI_STEP_NAME"""'], {}), False, 'import os\n'), (132, 'os.environ.get', 'os.environ.get', (['"""AI_EXPERIMENT_NAME"""'], {}), False, 'import os\n'), (142, 'os.environ.get', 'os.environ.get', (['"""MY_ENV"""'], {}), False, 'import os\n'), (181, 'os.path.exists', 'os.path.exists', (['file_path'], {}), False, 'import os\n'), (182, 'time.sleep', 'time.sleep', (['(1)'], {}), False, 'import time\n'), (51, 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', ([], {'field_path': '"""metadata.name"""'}), False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), (69, 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', ([], {'field_path': 'f"""metadata.annotations[\'{annotation}\']"""'}), False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), (87, 'kubernetes.client.V1ObjectFieldSelector', 'V1ObjectFieldSelector', ([], {'field_path': 'f"""metadata.labels[\'{label}\']"""'}), False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n'), (26, 'kubernetes.client.V1ResourceFieldSelector', 'V1ResourceFieldSelector', ([], {'container_name': '"""main"""', 'resource': 'resource', 'divisor': "('1m' if 'cpu' in resource else '1')"}), False, 'from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector, V1ResourceFieldSelector\n')]