#
# Generated by Elyra {{ elyra_version }}
#
import kfp
import kfp.aws
from kubernetes.client import *
from kubernetes.client.models import *

{# 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.image_pull_policy %}
    {{ task_name }}.container.set_image_pull_policy("{{ workflow_task.task_modifiers.image_pull_policy }}")
{%  endif %}
{%  if workflow_task.task_modifiers.special_output_files %}
    {{ task_name }}.output_artifact_paths["mlpipeline-metrics"] = "{{ workflow_task.task_modifiers.special_output_files.mlpipeline_metrics }}"
    {{ task_name }}.output_artifact_paths["mlpipeline-ui-metadata"] = "{{ workflow_task.task_modifiers.special_output_files.mlpipeline_ui_metadata }}"
{% endif %}
{%  if workflow_task.task_modifiers.object_storage_secret %}
    {{ task_name }}.apply(kfp.aws.use_aws_secret("{{ workflow_task.task_modifiers.object_storage_secret }}"))
{%  endif %}
    {{ task_name }}.set_display_name("{{ workflow_task.name | string_delimiter_safe }}")
{%  if workflow_task.doc %}
    {{ task_name }}.add_pod_annotation("elyra/node-user-doc","""{{ workflow_task.doc| string_delimiter_safe }}""")
{%  endif %}
{%  if workflow_task.task_modifiers.cpu_request %}
    {{ task_name }}.container.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 }}.container.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 }}.container.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 }}.container.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 }}.container.add_resource_limit(resource_name="{{ workflow_task.task_modifiers.gpu_limit.vendor }}", value="{{ workflow_task.task_modifiers.gpu_limit.size }}")
{%  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 }}.add_env_variable(V1EnvVar(name="{{ env_var_name }}", value="{{ env_var_value | string_delimiter_safe }}"))
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.set_run_name %}
{%   if workflow_engine == "tekton" %}
    {{ task_name }}.add_env_variable(V1EnvVar(name="ELYRA_RUN_NAME", value_from=V1EnvVarSource(field_ref=V1ObjectFieldSelector(field_path="metadata.annotations['pipelines.kubeflow.org/run_name']"))))
{%   else %}
    {{ task_name }}.add_env_variable(V1EnvVar(name="ELYRA_RUN_NAME", value="{{ workflow_task.task_modifiers.set_run_name }}"))
{%   endif %}
{%  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() %}
    {{ task_name }}.add_pod_label("{{ 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() %}
    {{ task_name }}.add_pod_annotation("{{ 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() %}
    {{ task_name }}.container.add_env_variable(V1EnvVar(
            name="{{ env_var }}",
            value_from=V1EnvVarSource(secret_key_ref=V1SecretKeySelector(name="{{ secret_dict.name }}", key="{{ secret_dict.key }}")),
        ))
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.kubernetes_volumes %}
{%    for volume_path, volume_dict in workflow_task.task_modifiers.kubernetes_volumes.items() %}
    {{ task_name }}.add_volume(
        V1Volume(
            name="{{ volume_dict.pvc_name}}",
            persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(claim_name="{{ volume_dict.pvc_name }}",),
        ))
    {{ task_name }}.container.add_volume_mount(
        V1VolumeMount(
            mount_path="{{ volume_path }}",
            name="{{ volume_dict.pvc_name }}",
{%      if volume_dict.sub_path %}            
            sub_path="{{ volume_dict.sub_path }}",
{%      endif %}
            read_only={{ volume_dict.read_only }},
        ))
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.kubernetes_tolerations %}
{%    for toleration_dict in workflow_task.task_modifiers.kubernetes_tolerations.values() %}
    {{ task_name }}.add_toleration(
        V1Toleration(
{%      if toleration_dict.effect %}
            effect="{{ toleration_dict.effect }}",
{%      else %}
            effect=None,
{%      endif %}
{%      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 %}
        ))
{%    endfor %}
{%  endif %}
{%  if workflow_task.task_modifiers.kubernetes_shared_mem_size %}
    {{ task_name }}.add_volume(V1Volume(
            name="shm",
            empty_dir=V1EmptyDirVolumeSource(medium="Memory", size_limit="{{ workflow_task.task_modifiers.kubernetes_shared_mem_size.size }}{{ workflow_task.task_modifiers.kubernetes_shared_mem_size.units }}"),
        ))
    {{ task_name }}.container.add_volume_mount(V1VolumeMount(mount_path="/dev/shm", name="shm"))
{%  endif %}
{%  if workflow_task.task_modifiers.crio_runtime %}
    {{ task_name }}.add_volume(V1Volume(
            name="{{ workflow_task.task_modifiers.crio_runtime.emptydir_volume_name }}",
            empty_dir=V1EmptyDirVolumeSource(medium="{{ workflow_task.task_modifiers.crio_runtime.emptydir_volume_medium }}", size_limit="{{ workflow_task.task_modifiers.crio_runtime.emptydir_volume_size }}"),
        ))
    {{ task_name }}.container.add_volume_mount(V1VolumeMount(mount_path="{{ workflow_task.task_modifiers.crio_runtime.emptydir_mount_path }}", name="{{ workflow_task.task_modifiers.crio_runtime.emptydir_volume_name }}"))
{%  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
{% if workflow_engine.lower() ==  "tekton" %}
    from kfp_tekton import compiler

    compiler.TektonCompiler().compile(
{% else %}
    kfp.compiler.Compiler().compile(
{% endif %}
        pipeline_func=generated_pipeline,
        package_path=Path(__file__).with_suffix(".yaml").name,
    )