from __future__ import absolute_import
from flytekit.models import common as _common
from flytekit.models.core import compiler as _compiler_models, identifier as _identifier, workflow as _core_workflow
from flyteidl.admin import workflow_pb2 as _admin_workflow


class WorkflowSpec(_common.FlyteIdlEntity):

    def __init__(self, template, sub_workflows):
        """
        This object fully encapsulates the specification of a workflow
        :param flytekit.models.core.workflow.WorkflowTemplate template:
        :param list[flytekit.models.core.workflow.WorkflowTemplate] sub_workflows:
        """
        self._template = template
        self._sub_workflows = sub_workflows

    @property
    def template(self):
        """
        :rtype: flytekit.models.core.workflow.WorkflowTemplate.WorkflowTemplate
        """
        return self._template

    @property
    def sub_workflows(self):
        """
        :rtype: list[flytekit.models.core.workflow.WorkflowTemplate.WorkflowTemplate]
        """
        return self._sub_workflows

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.admin.workflow_pb2.WorkflowSpec
        """
        return _admin_workflow.WorkflowSpec(
            template=self._template.to_flyte_idl(),
            sub_workflows=[s.to_flyte_idl() for s in self._sub_workflows],
        )

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param pb2_object: flyteidl.admin.workflow_pb2.WorkflowSpec
        :rtype: WorkflowSpec
        """
        return cls(_core_workflow.WorkflowTemplate.from_flyte_idl(pb2_object.template),
                   [_core_workflow.WorkflowTemplate.from_flyte_idl(s) for s in pb2_object.sub_workflows])


class Workflow(_common.FlyteIdlEntity):

    def __init__(
        self,
        id,
        closure
    ):
        """
        :param flytekit.models.core.identifier.Identifier id:
        :param WorkflowClosure closure:
        """
        self._id = id
        self._closure = closure

    @property
    def id(self):
        """
        :rtype: flytekit.models.core.identifier.Identifier
        """
        return self._id

    @property
    def closure(self):
        """
        :rtype: WorkflowClosure
        """
        return self._closure

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.admin.workflow_pb2.Workflow
        """
        return _admin_workflow.Workflow(
            id=self.id.to_flyte_idl(),
            closure=self.closure.to_flyte_idl()
        )

    @classmethod
    def from_flyte_idl(cls, pb2_object):
        """
        :param flyteidl.admin.workflow_pb2.Workflow pb2_object:
        :return: Workflow
        """
        return cls(
            id=_identifier.Identifier.from_flyte_idl(pb2_object.id),
            closure=WorkflowClosure.from_flyte_idl(pb2_object.closure)
        )


class WorkflowClosure(_common.FlyteIdlEntity):

    def __init__(self, compiled_workflow):
        """
        :param flytekit.models.core.compiler.CompiledWorkflowClosure compiled_workflow:
        """
        self._compiled_workflow = compiled_workflow

    @property
    def compiled_workflow(self):
        """
        :rtype: flytekit.models.core.compiler.CompiledWorkflowClosure
        """
        return self._compiled_workflow

    def to_flyte_idl(self):
        """
        :rtype: flyteidl.admin.workflow_pb2.WorkflowClosure
        """
        return _admin_workflow.WorkflowClosure(
            compiled_workflow=self.compiled_workflow.to_flyte_idl()
        )

    @classmethod
    def from_flyte_idl(cls, p):
        """
        :param flyteidl.admin.workflow_pb2.WorkflowClosure p:
        :rtype: WorkflowClosure
        """
        return cls(
            compiled_workflow=_compiler_models.CompiledWorkflowClosure.from_flyte_idl(p.compiled_workflow)
        )
