# Compile and convert kubeflow component
import base64
import re

from kfp.compiler._op_to_template import _op_to_template
from kfp.components import load_component_from_file
from kfp.dsl import PipelineParam, ContainerOp
from kfp.dsl._component_bridge import _create_container_op_from_component_and_arguments

valid_name_regex = r'^[A-Za-z][A-Za-z0-9_-]*$'


# KFP compile code will generate redudant parameter due to its original design
# In KFP's _data_passing_rewriter.py, it checks and fix. The function below did
# similar thing but with a different rule. It removes the redudant outputs parameter
def fix_parameter(template):
    """Fix the incorrect parameter generated by kfp compile"""
    if 'outputs' in template:
        parameters = template['outputs'].get('parameters')
        artifacts = template['outputs'].get('artifacts')
        if parameters is not None:
            temp = []
            for p in parameters:
                if artifacts is None or len(list(filter(lambda x: x.get('name') == p.get('name'), artifacts))) == 0:
                    temp.append(p)
            if len(temp) == 0:
                del template['outputs']['parameters']
            else:
                template['outputs']['parameters'] = temp
    return template


def validTemplateName(name: str):
    """Validate if the template name is compliant to the requirement"""
    return re.match(valid_name_regex, name)


def encode_k8s_name(input: str) -> str:
    """Encode name to be compatible with Kubernetes name convention"""
    return base64.b32encode(input.encode('utf-8')).decode('utf-8').lower().rstrip('=')


def restore_human_names(container_op: ContainerOp):
    """Remove the hash in the name"""
    container_op.name = container_op.human_name if validTemplateName(container_op.human_name) else encode_k8s_name(
        container_op.human_name)
    if container_op.output:
        container_op.output = PipelineParam(name=container_op.output.name, value=container_op.output.value)

    for k, v in container_op.outputs.items():
        container_op.outputs[k] = PipelineParam(name=container_op.output.name, value=container_op.output.value)


def component_to_template(file_name: str):
    """Convert to a component to an argo template"""
    component = load_component_from_file(file_name).component_spec
    args = {}
    for input in component.inputs:
        args[input.name] = PipelineParam(name=input.name, value="{{inputs.parameters." + input.name + "}}")
    container_op = _create_container_op_from_component_and_arguments(component, args)
    # Component to op process creating some hash name for tasks. Restore them here
    restore_human_names(container_op)
    return fix_parameter(_op_to_template(container_op))


def containerOp_to_template(op: ContainerOp):
    """Convert a container op to an argo template"""
    template = _op_to_template(op)
    template['name'] = op.human_name if validTemplateName(op.human_name) else encode_k8s_name(op.human_name)
    return fix_parameter(template)


def create_workflow_template_from_component_spec(file_name: str):
    """Create an argo workflowtemplate object from the given kfp component"""
    template = component_to_template(file_name)
    workflow_template = {
        'apiVersion': 'argoproj.io/v1alpha1',
        'kind': 'WorkflowTemplate',
        'metadata': {
            'name': encode_k8s_name(template['name']),
        },
        'spec': {
            'templates': [
                template
            ]
        }
    }
    return workflow_template


def create_workflow_template_from_container_op(op: ContainerOp):
    """Create an argo workflowtemplate object from a defined ContainerOp"""
    template = containerOp_to_template(op)
    workflow_template = {
        'apiVersion': 'argoproj.io/v1alpha1',
        'kind': 'WorkflowTemplate',
        'metadata': {
            'name': encode_k8s_name(template['name']),
        },
        'spec': {
            'templates': [
                template
            ]
        }
    }
    return workflow_template
