import kfp
{% if engine == "Tekton" %}
import kfp_tekton
{% endif %}
{% if kf_secured %}
import requests
import sys
{% endif %}
{% if cos_secret %}
from kfp.aws import use_aws_secret
{% endif %}
from datetime import datetime
from elyra.kfp.operator import ExecuteFileOp


@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 }} = ExecuteFileOp(name='{{ operation_id }}',
                                                      pipeline_name='{{ pipeline_name }}',
                                                      experiment_name='{{ experiment_name }}',
                                                      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_version='',
                                                      pipeline_source='{{ operation.pipeline_source }}',
                                                      pipeline_inputs={{ operation.pipeline_inputs }},
                                                      pipeline_outputs={{ operation.pipeline_outputs }},
        {% if operation.cpu_request %}
                                                      cpu_request='{{ operation.cpu_request }}',
        {% endif %}
        {% if operation.mem_request %}
                                                      mem_request='{{ operation.mem_request }}',
        {% endif %}
        {% if operation.gpu_limit %}
                                                      gpu_limit='{{ operation.gpu_limit }}',
        {% endif %}
                                                      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 }}
        {% if operation.container.image_pull_policy %}
            notebook_op_{{ loop.index }}.container.image_pull_policy = '{{ operation.container.image_pull_policy }}'
        {% endif %}
        {% if cos_secret %}
            ## Ensure that the secret named '{{ cos_secret }}' below is defined in the Kubernetes namespace '{{ namespace }}' where this pipeline will be run
            notebook_op_{{ loop.index }}.apply(use_aws_secret('{{ cos_secret }}'))
        {% endif %}

    {% endfor %}

{% if kf_secured %}
def get_istio_auth_session(url: str, username: str, password: str) -> dict:
    """
    Determine if the specified URL is secured by Dex and try to obtain a session cookie.
    WARNING: only Dex `staticPasswords` authentication is currently supported

    :param url: Kubeflow server URL, including protocol
    :param username: Dex `staticPasswords` username
    :param password: Dex `staticPasswords` password
    :return: auth session information
    """
    # define the default return object
    auth_session = {
        "endpoint_url": url,    # KF endpoint URL
        "redirect_url": None,   # KF redirect URL, if applicable
        "dex_login_url": None,  # Dex login URL (for POST of credentials)
        "is_secured": None,     # True if KF endpoint is secured
        "session_cookie": None  # Resulting session cookies in the form "key1=value1; key2=value2"
    }

    # use a persistent session (for cookies)
    with requests.Session() as s:

        ################
        # Determine if Endpoint is Secured
        ################
        resp = s.get(url, allow_redirects=True)
        if resp.status_code != 200:
            raise RuntimeError(
                f"HTTP status code '{resp.status_code}' for GET against: {url}"
            )

        auth_session["redirect_url"] = resp.url

        # if we were NOT redirected, then the endpoint is UNSECURED
        if len(resp.history) == 0:
            auth_session["is_secured"] = False
            return auth_session
        else:
            auth_session["is_secured"] = True

        ################
        # Get Dex Login URL (that allows us to POST credentials)
        ################
        redirect_url_obj = urlsplit(auth_session["redirect_url"])

        # we expect a "Dex Login" URL to have `/auth` in the HTTP path
        if "/auth" not in redirect_url_obj.path:
            raise RuntimeError(
                f"Path of `redirect_url` must contain '/auth', but got: {auth_session['redirect_url']}"
            )

        # rewrite "Dex Login" URL to "/auth/local" (for dex `staticPasswords` login)
        # (needed when dex is configured with multiple auth options)
        if "/auth/local" not in redirect_url_obj.path:
            redirect_url_obj = redirect_url_obj._replace(
                path=redirect_url_obj.path.replace("/auth", "/auth/local")
            )
            auth_session["redirect_url"] = redirect_url_obj.geturl()

        # Get redirected to page we can POST credentials to
        # (needed when dex is configured with multiple auth options)
        resp = s.get(auth_session["redirect_url"], allow_redirects=True)
        if resp.status_code != 200:
            raise RuntimeError(
                f"HTTP status code '{resp.status_code}' for GET against: {auth_session['redirect_url']}"
            )

        auth_session["dex_login_url"] = resp.url

        ################
        # Attempt Dex `staticPasswords` Login
        ################
        resp = s.post(
            auth_session["dex_login_url"],
            data={"login": username, "password": password},
            allow_redirects=True
        )
        if len(resp.history) == 0:
            raise RuntimeError(
                f"No redirect after POST to: {auth_session['dex_login_url']} - "
                f"Login credentials were probably invalid."
            )

        # store the session cookies in a "key1=value1; key2=value2" string
        auth_session["session_cookie"] = "; ".join([f"{c.name}={c.value}" for c in s.cookies])

    return auth_session
{% endif %}


if __name__ == "__main__":

{% if kf_secured %}
    # user id and password are required parameters
    if(len(sys.argv) < 3):
        print(f'Invocation: python {sys.argv[0]} '
               '<kf-user-id> <kf-password>')
        sys.exit(1)

    print('Trying to authenticate with Kubeflow server ...')
    api_endpoint = '{{ api_endpoint }}'.rstrip('/')

    try:
        # attempt to create a session cookies using the provided `api_username` and `api_password`
        auth_session = get_istio_auth_session(
            url=api_endpoint,
            username=api_username,
            password=api_password
        )
    except Exception as ex:
        raise RuntimeError(
            f"Failed to create istio auth session for Kubeflow endpoint: '{api_endpoint}'"
        ) from ex

    auth_cookie = auth_session['session_cookie']
{% else %}
    auth_cookie = None
{% endif %}

    pipeline_func = create_pipeline
    pipeline_name = '{{ pipeline_name }}'
    pipeline_archive = pipeline_name + '.pipeline.tgz'
    experiment_name = '{{ experiment_name }}'

    timestamp = datetime.now().strftime("%m%d%H%M%S")

    print('Compiling pipeline using {{ engine }} compiler ...')
{% if engine == "Tekton" %}
    kfp_tekton.compiler.TektonCompiler().compile(pipeline_func, pipeline_archive)

    client = kfp_tekton.TektonClient(host='{{ api_endpoint }}',
                                     cookies=auth_cookie)
{% else %}
    kfp.compiler.Compiler().compile(pipeline_func, pipeline_archive)

    client = kfp.Client(host='{{ api_endpoint }}',
                        cookies=auth_cookie)
{% endif %}

    pipeline_id = client.get_pipeline_id(pipeline_name)
    if pipeline_id is None:
        # Upload new pipeline. The call returns a unique pipeline id.
        print(f'Uploading pipeline {pipeline_name} ...')
        kfp_pipeline = \
            client.upload_pipeline(pipeline_archive,
                                   pipeline_name,
                                   '{{ pipeline_description }}')
        pipeline_id = kfp_pipeline.id
        version_id = None     
    else:
        # Append timestamp to generate unique version name
        pipeline_version_name = f'{pipeline_name}-{timestamp}'
        # Upload a pipeline version. The call returns a unique version id.
        print(f'Uploading pipeline version {pipeline_version_name} ...')
        kfp_pipeline = \
            client.upload_pipeline_version(pipeline_archive,
                                           pipeline_version_name,
                                           pipeline_id=pipeline_id)
        version_id = kfp_pipeline.id

    # Create an experiment
    print(f'Creating experiment {experiment_name} ...')
    experiment = client.create_experiment(experiment_name,
    {% if namespace is none  %}
                                          namespace=None)
    {% else %}
                                          namespace='{{ namespace }}')
    {% endif %}

    # Generate unique identifier for the pipeline run
    run_name = f'{pipeline_name}-{timestamp}'

    # Run the pipeline (or specified pipeline version)
    print(f'Starting pipeline run {run_name} ...')
    run = client.run_pipeline(experiment_id=experiment.id,
                              job_name=run_name,
                              pipeline_id=pipeline_id,
                              version_id=version_id)

    print('Pipeline run monitoring URL: '
          f'{{ api_endpoint }}/#/runs/details/{run.id}')
