#
# Generated by Elyra {{ elyra_version }}
#
import kfp
from kubernetes.client import *
from kubernetes.client.models import *
from kfp.kubernetes import secret, volume

from typing import Optional

# ------------------------------------------------------------------
# kfp-kubernetes 1.1.0 is misisng these function, explicity using them,
# TODO: remove these function once a new release of kfp-kubernetes is made.
# ------------------------------------------------------------------

from google.protobuf import json_format
from kfp.dsl import PipelineTask
from kfp.kubernetes import common
from kfp.kubernetes import kubernetes_executor_config_pb2 as pb

from kfp.kubernetes import add_toleration, add_pod_label, add_pod_annotation

{# Load statements for custom components                            -#}
{# component_hash = """<TEXT>"""                                  -#}
{# factory_hash = kfp.components.load_component_from_text(component_hash) -#}
{% for hash, component_definition in component_definitions.items() %}
component_def_{{ hash | python_safe }} = """
{{ component_definition }}
"""

factory_{{ hash | python_safe }} = kfp.components.load_component_from_text(component_def_{{ hash | python_safe }})
{% endfor %}

{#  Define pipeline  -#}
{% if pipeline_description %}
@kfp.dsl.pipeline(name="{{ pipeline_name }}", description="{{ pipeline_description | string_delimiter_safe }}")
{% else %}
@kfp.dsl.pipeline(name="{{ pipeline_name }}")
{% endif %}
def generated_pipeline(
{% if pipeline_parameters %}
{% for parameter in pipeline_parameters %}
    {{ parameter.name }}{% if parameter.input_type.type_hint %}: {{ parameter.input_type.type_hint }}{% endif %} = {{ parameter|param_val_to_python_var }},
{% endfor %}
{% endif %}
):
{% for workflow_task in workflow_tasks.values() %}
    {% set task_name = "task_" + workflow_task.escaped_task_id %}
    # Task for node '{{ workflow_task.name }}'
    {{ task_name }} = factory_{{ workflow_task.component_definition_hash | python_safe }}(
{%  for task_input_name, task_input_spec in workflow_task.task_inputs.items() %}
{%     if task_input_spec.task_output_reference %}
        {{ task_input_name }}=task_{{ task_input_spec.task_output_reference.task_id }}.outputs["{{ task_input_spec.task_output_reference.output_id }}"],
{%     elif task_input_spec.pipeline_parameter_reference %}
        {{ task_input_name }}={{ task_input_spec.pipeline_parameter_reference }},
{%     elif task_input_spec.requires_quoted_rendering %}
        {{ task_input_name }}="""{{ task_input_spec.value | string_delimiter_safe }}""",
{%     else %}
        {{ task_input_name }}={{ task_input_spec.value }},
{%     endif %}
{%  endfor %}
    )
{%  if workflow_task.task_modifiers.object_storage_secret %}
    secret.use_secret_as_env({{ task_name }}, "{{ workflow_task.task_modifiers.object_storage_secret }}", { "AWS_ACCESS_KEY_ID": "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY": "AWS_SECRET_ACCESS_KEY" })
{%  endif %}
    {{ task_name }}.set_display_name("{{ workflow_task.name | string_delimiter_safe }}")
{%  if workflow_task.doc %}
    add_pod_annotation({{ task_name }}, "elyra/node-user-doc","""{{ workflow_task.doc| string_delimiter_safe }}""")
{%  endif %}
{%  if workflow_task.task_modifiers.cpu_request %}
    {{ task_name }}.set_cpu_request(cpu="{{ workflow_task.task_modifiers.cpu_request }}")
{%  endif %}
{%  if workflow_task.task_modifiers.mem_request and workflow_task.task_modifiers.mem_request.size %}
    {{ task_name }}.set_memory_request(memory="{{ workflow_task.task_modifiers.mem_request.size }}{{ workflow_task.task_modifiers.mem_request.units }}")
{%  endif %}
{%  if workflow_task.task_modifiers.cpu_limit %}
    {{ task_name }}.set_cpu_limit(cpu="{{ workflow_task.task_modifiers.cpu_limit }}")
{%  endif %}
{%  if workflow_task.task_modifiers.memory_limit and workflow_task.task_modifiers.memory_limit.size %}
    {{ task_name }}.set_memory_limit(memory="{{ workflow_task.task_modifiers.memory_limit.size }}{{ workflow_task.task_modifiers.memory_limit.units }}")
{%  endif %}
{%  if workflow_task.task_modifiers.gpu_limit and workflow_task.task_modifiers.gpu_limit.size %}
    {{ task_name }}.set_accelerator_limit("{{ workflow_task.task_modifiers.gpu_limit.size }}").set_accelerator_type("{{ workflow_task.task_modifiers.gpu_limit.vendor }}")
{%  endif %}
{%  if workflow_task.task_modifiers.env_variables %}
{%    for env_var_name, env_var_value in workflow_task.task_modifiers.env_variables.items() %}
    {{ task_name }}.set_env_variable(name="{{ env_var_name }}", value="{{ env_var_value | string_delimiter_safe }}")
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.set_run_name %}
    {{ task_name }}.set_env_variable(name="ELYRA_RUN_NAME", value="{{ workflow_task.task_modifiers.set_run_name }}")
{%  endif %}
{%  if workflow_task.task_modifiers.disable_node_caching %}
    {{ task_name }}.execution_options.caching_strategy.max_cache_staleness = "P0D"
{%  endif %}
{%  if workflow_task.task_modifiers.pod_labels %}
{%    for pod_label_key, pod_label_value in workflow_task.task_modifiers.pod_labels.items() %}
    add_pod_label({{ task_name }}, "{{ pod_label_key }}", "{{ pod_label_value }}")
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.pod_annotations %}
{%    for pod_annotation_key, pod_annotation_value in workflow_task.task_modifiers.pod_annotations.items() %}
    add_pod_annotation({{ task_name }}, "{{ pod_annotation_key }}" , """{{ pod_annotation_value | string_delimiter_safe }}""")
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.kubernetes_secrets %}
{%    for env_var, secret_dict in workflow_task.task_modifiers.kubernetes_secrets.items() %}
    secret.use_secret_as_env({{ task_name }}, "{{ secret_dict.name }}", { "{{ secret_dict.key }}" : "{{ env_var }}" })
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.kubernetes_volumes %}
{%    for volume_path, volume_dict in workflow_task.task_modifiers.kubernetes_volumes.items() %}
    volume.mount_pvc({{ task_name }}, "{{ volume_dict.pvc_name }}", "{{ volume_path }}")
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.kubernetes_tolerations %}
{% for toleration_dict in workflow_task.task_modifiers.kubernetes_tolerations.values() %}
    add_toleration(
        {{ task_name }},
        {% if toleration_dict.key %}
            key="{{ toleration_dict.key }}",
        {% else %}
            key=None,
        {% endif %}
        operator="{{ toleration_dict.operator }}",
        {% if toleration_dict.value %}
            value="{{ toleration_dict.value | string_delimiter_safe }}",
        {% else %}
            value=None,
        {% endif %}
        {% if toleration_dict.effect %}
            effect="{{ toleration_dict.effect }}",
        {% else %}
            effect=None,
        {% endif %}
    )
{% endfor %}
{%  endif %}
{#  declare upstream dependencies -#}
{%  if workflow_task.upstream_workflow_task_ids %}
{%    for upstream_workflow_task_id in workflow_task.upstream_workflow_task_ids %}
    {{ task_name }}.after(task_{{ upstream_workflow_task_id | python_safe }})
{%    endfor %}
{%  endif %}
{% endfor %}

if __name__ == "__main__":
    from pathlib import Path
    kfp.compiler.Compiler().compile(
        pipeline_func=generated_pipeline,
        package_path=Path(__file__).with_suffix(".yaml").name,
    )