# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import hashlib
import os
import sys
import kubernetes
import yaml
from time import sleep

from metadata_helpers import *


namespace_to_watch = os.environ.get('NAMESPACE_TO_WATCH', 'default')


kubernetes.config.load_incluster_config()
k8s_api = kubernetes.client.CoreV1Api()
k8s_watch = kubernetes.watch.Watch()


patch_retries = 20
sleep_time = 0.1


def patch_pod_metadata(
    namespace: str,
    pod_name: str,
    patch: dict,
    k8s_api: kubernetes.client.CoreV1Api = None,
):
    k8s_api = k8s_api or kubernetes.client.CoreV1Api()
    patch = {
        'metadata': patch
    }
    for retry in range(patch_retries):
        try:
            pod = k8s_api.patch_namespaced_pod(
                name=pod_name,
                namespace=namespace,
                body=patch,
            )
            return pod
        except Exception as e:
            print(e)
            sleep(sleep_time)


#Connecting to MetadataDB
mlmd_store = connect_to_mlmd()
print("Connected to the metadata store")

PIPELINE_RUNTIME = os.getenv("PIPELINE_RUNTIME", "argo").lower()

ARGO_OUTPUTS_ANNOTATION_KEY = 'workflows.argoproj.io/outputs'
ARGO_TEMPLATE_ANNOTATION_KEY = 'workflows.argoproj.io/template'
KFP_COMPONENT_SPEC_ANNOTATION_KEY = 'pipelines.kubeflow.org/component_spec'
METADATA_EXECUTION_ID_LABEL_KEY = 'pipelines.kubeflow.org/metadata_execution_id'
METADATA_CONTEXT_ID_LABEL_KEY = 'pipelines.kubeflow.org/metadata_context_id'
METADATA_ARTIFACT_IDS_ANNOTATION_KEY = 'pipelines.kubeflow.org/metadata_artifact_ids'
METADATA_INPUT_ARTIFACT_IDS_ANNOTATION_KEY = 'pipelines.kubeflow.org/metadata_input_artifact_ids'
METADATA_OUTPUT_ARTIFACT_IDS_ANNOTATION_KEY = 'pipelines.kubeflow.org/metadata_output_artifact_ids'

ARGO_WORKFLOW_LABEL_KEY = 'workflows.argoproj.io/workflow'
ARGO_COMPLETED_LABEL_KEY = 'workflows.argoproj.io/completed'
METADATA_WRITTEN_LABEL_KEY = 'pipelines.kubeflow.org/metadata_written'

TEKTON_PIPELINERUN_LABEL_KEY = 'tekton.dev/pipelineRun'
TEKTON_READY_ANNOTATION_KEY = 'tekton.dev/ready'
TEKTON_TASKRUN_LABEL_KEY = 'tekton.dev/taskRun'
TEKTON_PIPELINETASK_LABEL_KEY = 'tekton.dev/pipelineTask'
TEKTON_INPUT_ARTIFACT_ANNOTATION_KEY = 'tekton.dev/input_artifacts'
TEKTON_OUTPUT_ARTIFACT_ANNOTATION_KEY = 'tekton.dev/output_artifacts'

PIPELINE_LABEL_KEY = TEKTON_PIPELINERUN_LABEL_KEY if PIPELINE_RUNTIME == "tekton" else ARGO_WORKFLOW_LABEL_KEY


def output_name_to_argo(name: str) -> str:
    import re
    return re.sub('-+', '-', re.sub('[^-0-9a-z]+', '-', name.lower())).strip('-')


def artifact_to_uri(artifact: dict) -> str:
    if 's3' in artifact:
        s3_artifact = artifact['s3']
        return 'minio://{bucket}/{key}'.format(
            bucket=s3_artifact.get('bucket', ''),
            key=s3_artifact.get('key', ''),
        )
    elif 'raw' in artifact:
        return None
    else:
        return None


def is_tfx_pod(pod) -> bool:
    main_step_name = 'step-main' if PIPELINE_RUNTIME == "tekton" else 'main'
    main_containers = [container for container in pod.spec.containers if container.name == main_step_name]
    if len(main_containers) != 1:
        return False
    main_container = main_containers[0]
    return main_container.command and main_container.command[-1].endswith('tfx/orchestration/kubeflow/container_entrypoint.py')


def get_component_template(obj):
    '''
    Return an Argo formatted component template
    '''
    if PIPELINE_RUNTIME == "tekton":
        artifacts = json.loads(obj.metadata.annotations[TEKTON_INPUT_ARTIFACT_ANNOTATION_KEY])
        results = artifacts.get(obj.metadata.labels[TEKTON_PIPELINETASK_LABEL_KEY], [])
        component_template = {
            "name": obj.metadata.labels[TEKTON_PIPELINETASK_LABEL_KEY],
            "inputs": {
                "artifacts": [
                    {
                        "name": i["name"],
                        "s3": {
                            "bucket": "mlpipeline",
                            "key": "artifacts/%s/%s/%s.tgz" % (obj.metadata.labels[TEKTON_PIPELINERUN_LABEL_KEY],
                                                               i['parent_task'],
                                                               i["name"].replace(i['parent_task'] + '-', ''))
                        }
                    } for i in results
                ]
            }
        }
        return component_template
    else:
        return json.loads(obj.metadata.annotations[ARGO_TEMPLATE_ANNOTATION_KEY])


def get_output_template(obj):
    '''
    Return an Argo formatted artifact output template
    '''
    if PIPELINE_RUNTIME == "tekton":
        artifacts = json.loads(obj.metadata.annotations[TEKTON_OUTPUT_ARTIFACT_ANNOTATION_KEY])
        results = artifacts.get(obj.metadata.labels[TEKTON_PIPELINETASK_LABEL_KEY], [])
        artifact_prefix = obj.metadata.labels[TEKTON_PIPELINETASK_LABEL_KEY] + '-'
        s3_key_prefix = 'artifacts/%s/%s' % (obj.metadata.labels[TEKTON_PIPELINERUN_LABEL_KEY],
                                             obj.metadata.labels[TEKTON_PIPELINETASK_LABEL_KEY])
        output_template = {
            "name": obj.metadata.labels[TEKTON_PIPELINETASK_LABEL_KEY],
            "artifacts": [
                {
                    "name": i["name"],
                    "path": i["path"],
                    "s3": {
                        "bucket": "mlpipeline",
                        "key": "%s/%s.tgz" % (s3_key_prefix,
                                              i["name"].replace(artifact_prefix, ''))
                    }
                } for i in results
            ]
        }
        return output_template
    else:
        return json.loads(obj.metadata.annotations[ARGO_OUTPUTS_ANNOTATION_KEY])


# Caches (not expected to be persistent)
# These caches are only used to prevent race conditions. Race conditions happen because the writer can see multiple versions of K8s object before the applied labels show up.
# They are expected to be lost when restarting the service.
# The operation of the Metadata Writer remains correct even if it's getting restarted frequently. (Kubernetes only sends the latest version of resource for new watchers.)
# Technically, we could remove the objects from cache as soon as we see that our labels have been applied successfully.
pod_name_to_execution_id = {}
workflow_name_to_context_id = {}
pods_with_written_metadata = set()

while True:
    print("Start watching Kubernetes Pods created by Argo")
    for event in k8s_watch.stream(
        k8s_api.list_namespaced_pod,
        namespace=namespace_to_watch,
        label_selector=PIPELINE_LABEL_KEY,
    ):
        try:
            obj = event['object']
            print('Kubernetes Pod event: ', event['type'], obj.metadata.name, obj.metadata.resource_version)
            if event['type'] == 'ERROR':
                print(event)

            pod_name = obj.metadata.name

            # Logging pod changes for debugging
            with open('/tmp/pod_' + obj.metadata.name + '_' + obj.metadata.resource_version, 'w') as f:
                f.write(yaml.dump(obj.to_dict()))

            assert obj.kind == 'Pod'

            if METADATA_WRITTEN_LABEL_KEY in obj.metadata.labels:
                continue

            # Skip TFX pods - they have their own metadata writers
            if is_tfx_pod(obj):
                continue

            pipeline_name = obj.metadata.labels[PIPELINE_LABEL_KEY] # Should exist due to initial filtering
            template = get_component_template(obj)
            template_name = template['name']

            component_name = template_name
            component_version = component_name
            output_name_to_type = {}
            if KFP_COMPONENT_SPEC_ANNOTATION_KEY in obj.metadata.annotations:
                component_spec_text = obj.metadata.annotations[KFP_COMPONENT_SPEC_ANNOTATION_KEY]
                component_spec = json.loads(component_spec_text)
                component_spec_digest = hashlib.sha256(component_spec_text.encode()).hexdigest()
                component_name = component_spec.get('name', component_name)
                component_version = component_name + '@sha256=' + component_spec_digest
                output_name_to_type = {output['name']: output.get('type', None) for output in component_spec.get('outputs', [])}
                output_name_to_type = {output_name_to_argo(k): v for k, v in output_name_to_type.items() if v}

            if obj.metadata.name in pod_name_to_execution_id:
                execution_id = pod_name_to_execution_id[obj.metadata.name]
                context_id = workflow_name_to_context_id[pipeline_name]
            elif METADATA_EXECUTION_ID_LABEL_KEY in obj.metadata.labels:
                execution_id = int(obj.metadata.labels[METADATA_EXECUTION_ID_LABEL_KEY])
                context_id = int(obj.metadata.labels[METADATA_CONTEXT_ID_LABEL_KEY])
                print('Found execution id: {}, context id: {} for pod {}.'.format(execution_id, context_id, obj.metadata.name))
            else:
                run_context = get_or_create_run_context(
                    store=mlmd_store,
                    run_id=pipeline_name, # We can switch to internal run IDs once backend starts adding them
                )

                # Adding new execution to the database
                execution = create_new_execution_in_existing_run_context(
                    store=mlmd_store,
                    context_id=run_context.id,
                    execution_type_name=KFP_EXECUTION_TYPE_NAME_PREFIX + component_version,
                    pod_name=pod_name,
                    pipeline_name=pipeline_name,
                    run_id=pipeline_name,
                    instance_id=component_name,
                )

                input_artifacts = template.get('inputs', {}).get('artifacts', [])
                input_artifact_ids = []
                for input_artifact in input_artifacts:
                    artifact_uri = artifact_to_uri(input_artifact)
                    if not artifact_uri:
                        continue

                    if PIPELINE_RUNTIME == "tekton":
                        input_name = input_artifact.get('name', '')
                        input_prefix = template_name + '-'
                        input_name = input_name[len(input_prefix):]
                    else:
                        input_name = input_artifact.get('path', '') # Every artifact should have a path in Argo
                        input_artifact_path_prefix = '/tmp/inputs/'
                        input_artifact_path_postfix = '/data'
                        if input_name.startswith(input_artifact_path_prefix):
                            input_name = input_name[len(input_artifact_path_prefix):]
                        if input_name.endswith(input_artifact_path_postfix):
                            input_name = input_name[0: -len(input_artifact_path_postfix)]

                    artifact = link_execution_to_input_artifact(
                        store=mlmd_store,
                        execution_id=execution.id,
                        uri=artifact_uri,
                        input_name=input_name,
                    )
                    if artifact is None:
                        # TODO: Maybe there is a better way to handle missing upstream artifacts
                        continue

                    input_artifact_ids.append(dict(
                        id=artifact.id,
                        name=input_name,
                        uri=artifact.uri,
                    ))
                    print('Found Input Artifact: ' + str(dict(
                        input_name=input_name,
                        id=artifact.id,
                        uri=artifact.uri,
                    )))

                execution_id = execution.id
                context_id = run_context.id

                obj.metadata.labels[METADATA_EXECUTION_ID_LABEL_KEY] = execution_id
                obj.metadata.labels[METADATA_CONTEXT_ID_LABEL_KEY] = context_id

                metadata_to_add = {
                    'labels': {
                        METADATA_EXECUTION_ID_LABEL_KEY: str(execution_id),
                        METADATA_CONTEXT_ID_LABEL_KEY: str(context_id),
                    },
                    'annotations': {
                        METADATA_INPUT_ARTIFACT_IDS_ANNOTATION_KEY: json.dumps(input_artifact_ids),
                    },
                }

                patch_pod_metadata(
                    namespace=obj.metadata.namespace,
                    pod_name=obj.metadata.name,
                    patch=metadata_to_add,
                )
                pod_name_to_execution_id[obj.metadata.name] = execution_id
                workflow_name_to_context_id[pipeline_name] = context_id

                print('New execution id: {}, context id: {} for pod {}.'.format(execution_id, context_id, obj.metadata.name))

                print('Execution: ' + str(dict(
                    context_id=context_id,
                    context_name=pipeline_name,
                    execution_id=execution_id,
                    execution_name=obj.metadata.name,
                    component_name=component_name,
                )))

                # TODO: Log input parameters as execution options.
                # Unfortunately, DSL compiler loses the information about inputs and their arguments.

            if (
                obj.metadata.name not in pods_with_written_metadata
                and (
                    obj.metadata.labels.get(ARGO_COMPLETED_LABEL_KEY, 'false') == 'true'
                    or ARGO_OUTPUTS_ANNOTATION_KEY in obj.metadata.annotations
                    or TEKTON_OUTPUT_ARTIFACT_ANNOTATION_KEY in obj.metadata.annotations
                )
            ):
                artifact_ids = []

                if ARGO_OUTPUTS_ANNOTATION_KEY in obj.metadata.annotations or TEKTON_OUTPUT_ARTIFACT_ANNOTATION_KEY in obj.metadata.annotations: # Should be present
                    outputs = get_output_template(obj)
                    pipeline_output_artifacts = {}

                    for artifact in outputs.get('artifacts', []):
                        art_name = artifact['name']
                        output_prefix = template_name + '-'
                        if art_name.startswith(output_prefix):
                            art_name = art_name[len(output_prefix):]
                        pipeline_output_artifacts[art_name] = artifact
                    
                    output_artifacts = []
                    for name, art in pipeline_output_artifacts.items():
                        artifact_uri = artifact_to_uri(art)
                        if not artifact_uri:
                            continue
                        artifact_type_name = output_name_to_type.get(name, 'NoType') # Cannot be None or ''

                        print('Adding Output Artifact: ' + str(dict(
                            output_name=name,
                            uri=artifact_uri,
                            type=artifact_type_name,
                        )))

                        artifact = create_new_output_artifact(
                            store=mlmd_store,
                            execution_id=execution_id,
                            context_id=context_id,
                            uri=artifact_uri,
                            type_name=artifact_type_name,
                            output_name=name,
                            #run_id='Context_' + str(context_id) + '_run',
                            run_id=pipeline_name,
                            argo_artifact=art,
                        )

                        artifact_ids.append(dict(
                            id=artifact.id,
                            name=name,
                            uri=artifact_uri,
                            type=artifact_type_name,
                        ))

                metadata_to_add = {
                    'labels': {
                        METADATA_WRITTEN_LABEL_KEY: 'true',
                    },
                    'annotations': {
                        METADATA_OUTPUT_ARTIFACT_IDS_ANNOTATION_KEY: json.dumps(artifact_ids),
                    },
                }
                
                patch_pod_metadata(
                    namespace=obj.metadata.namespace,
                    pod_name=obj.metadata.name,
                    patch=metadata_to_add,
                )

                pods_with_written_metadata.add(obj.metadata.name)

        except Exception as e:
            import traceback
            print(traceback.format_exc())
