import kfp
from kfp_notebook.pipeline import NotebookOp


@kfp.dsl.pipeline(
   name='{{ pipeline_name }}',
   description='{{ pipeline_description }}'
)

def create_pipeline():
    {% set metrics_file = writable_container_dir ~ '/mlpipeline-metrics.json' %}
    {% set metadata_file = writable_container_dir ~ '/mlpipeline-ui-metadata.json' %}
    {% for key, operation in operations_list.items() %}
        {% set op_name = operation.notebook_name.split('.') %}
        {% set operation_id = op_name[0] | to_basename %}
            # Operator for "{{ operation.notebook_name | to_basename }}"
            notebook_op_{{ loop.index }} = NotebookOp(name='{{ operation_id }}',
                                                      notebook='{{ operation.notebook_name | to_basename }}',
                                                      cos_endpoint='{{ operation.cos_endpoint }}',
                                                      cos_bucket='{{ operation.cos_bucket }}',
                                                      cos_directory='{{ operation.cos_directory }}',
                                                      cos_dependencies_archive='{{ operation.cos_dependencies_archive }}',
                                                      pipeline_inputs={{ operation.pipeline_inputs }},
                                                      pipeline_outputs={{ operation.pipeline_outputs }},
                                                      image='{{ operation.image }}',
                                                      file_outputs={
                                                        'mlpipeline-metrics': '{{ metrics_file }}',
                                                        'mlpipeline-ui-metadata': '{{ metadata_file }}'})

            notebook_op_{{ loop.index }}.name = '{{ operation.name }}'
            notebook_op_{{ loop.index }}.env_variables = {{ operation.env_variables }}
            notebook_op_{{ loop.index }}.dependent_names = {{ operation.dependent_names }}

    {% endfor %}


if __name__ == "__main__":
    pipeline_func = create_pipeline
    pipeline_filename = '{{ pipeline_name }}' + '.pipeline.tgz'
    kfp.compiler.Compiler().compile(pipeline_func, pipeline_filename)

    client = kfp.Client(host='{{ api_endpoint }}')
    name = '{{ pipeline_name }}'
    exp = client.create_experiment('{{ pipeline_name }}')
    run_result = client.run_pipeline(exp.id, name, pipeline_filename, {})

    print("Pipeline Submitted: {{ api_endpoint }}/#/runs/details/" + run_result.id)
