from flytekit import task, workflow
from flytekit.remote import FlyteRemote
from flytekit.configuration import Config

remote = FlyteRemote(
    Config.auto(), 
    default_project="flytesnacks", 
    default_domain="development", 
    interactive_mode_enabled=True,
)

# Define a task that produces the string "Hello, World!"
# by using the `@task` decorator to annotate the Python function
@task
def say_hello() -> str:
    return "Hello, World!"


# Handle the output of a task like that of a regular Python function.
@workflow
def hello_world_wf() -> str:
    res = say_hello()
    return res 


# Run the workflow locally by calling it like a Python function
if __name__ == "__main__":
    print(f"Running hello_world_wf() {hello_world_wf()}")

    # Execute the task
    exe = remote.execute(hello, inputs={"name": "Flyte"})

    # This will print the URL to the console
    print(exe.execution_url)

    # Wait for the task to complete
    exe = exe.wait(poll_interval=1)

    # Print the outputs
    print(exe.outputs)

class WorkflowCommand(click.RichGroup):
    
    # 根据文件名从远程或本地确定指定文件，如果远程下载文件
    def __init__(self, filename: str, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if ctx.file_access.is_remote(filename):
            local_path = os.path.join(os.path.curdir, filename.rsplit("/", 1)[1])
            ctx.file_access.download(filename, local_path)
            self._filename = pathlib.Path(local_path).resolve()
            self._should_delete = True
        else:
            self._filename = pathlib.Path(filename).resolve()
            self._should_delete = False

    def list_commands(self, ctx):
        # 通过文件解析task、wf、lp等装饰器实体方法
        entities = get_entities_in_file(self._filename, self._should_delete)
        self._entities = entities
        # 返回方法中的实体名称
        return entities.all()

    def get_command(self, ctx, exe_entity):
        # 先确定工作目录地址
        project_root = _find_project_root(self._filename)
        # 构造运行环境参数
        run_level_params: RunLevelParams = ctx.obj
        run_level_params.computed_params.project_root = project_root
        run_level_params.computed_params.module = module
        # 加载目标模块(exe_entity)依赖包和代码
        entity = load_naive_entity(module, exe_entity, project_root)
        return self._create_command(ctx, exe_entity, run_level_params, entity, entity_type)

     def _create_command(
        self,
        ctx: click.Context,
        entity_name: str,
        run_level_params: RunLevelParams,
        loaded_entity: [PythonTask, WorkflowBase, LaunchPlan],
        entity_type: str,
    ):
        # 获取remote实例(FlyteRemote)
        r = run_level_params.remote_instance()
        flyte_ctx = r.context

        final_inputs_with_defaults = loaded_entity.python_interface.inputs_with_defaults
        if isinstance(loaded_entity, LaunchPlan):
            # 从默认输入中提取默认值
            all_inputs = loaded_entity.python_interface.inputs_with_defaults
            default_inputs = loaded_entity.saved_inputs
            pmap = loaded_entity.parameters
            final_inputs_with_defaults = {}
            for name, _ in pmap.parameters.items():
                _type, v = all_inputs[name]
                if name in default_inputs:
                    v = default_inputs[name]
                final_inputs_with_defaults[name] = _type, v

        # Add options for each of the workflow inputs
        params = []
        for input_name, input_type_val in final_inputs_with_defaults.items():
            literal_var = loaded_entity.interface.inputs.get(input_name)
            python_type, default_val = input_type_val
            required = type(None) not in get_args(python_type) and default_val is None
            params.append(to_click_option(ctx, flyte_ctx, input_name, literal_var, python_type, default_val, required))

        h = f"{click.style(entity_type, bold=True)} ({run_level_params.computed_params.module}.{entity_name})"
        if isinstance(loaded_entity, LaunchPlan):
            h = h + click.style(f" (LP Name: {loaded_entity.name})", fg="yellow")
        else:
            if loaded_entity.__doc__:
                h = h + click.style(f" {loaded_entity.__doc__}", dim=True)
        # 同时也从yaml或json文件中读取参数
        # 最终把参数传递给run_command方法中
        cmd = YamlFileReadingCommand(
            name=entity_name,
            params=params,
            help=h,
            callback=run_command(ctx, loaded_entity),
        )
        return cmd

    def _run(*args, **kwargs):
        # ...
        # 这里有个有一些关键点, 如果命令行没有--remote参数，则在本地运行，如果是feature值，再通过asyncio进行异步运行
        if not run_level_params.is_remote:
            with FlyteContextManager.with_context(_update_flyte_context(run_level_params)):
                if run_level_params.envvars:
                    for env_var, value in run_level_params.envvars.items():
                        os.environ[env_var] = value
                if run_level_params.overwrite_cache:
                    os.environ["FLYTE_LOCAL_CACHE_OVERWRITE"] = "true"
                # 本地同步
                output = entity(**inputs)
                if inspect.iscoroutine(output):
                    # 如果是协程函数，异步执行
                    # TODO: make eager mode workflows run with local-mode
                    output = asyncio.run(output)
                click.echo(output)
                return

        # 以下是进行远程执行
        # 创建运行上下文，对这次运行所维护的状态信息进行入栈，FlyteContextManager稍后进行解析
        with context_manager.FlyteContextManager.with_context(remote.context.new_builder()):
            # 向后端服务器发送请求，创建相关entity数据库记录，并会将输入和输出存储远程对象存储，将相关的ref信息返回
            # 这其中包括：1.打包代码文件；2.上传至对象存储，并获得引用; 3.生成执行setting，序列号选项等
            remote_entity = remote.register_script(...)
            # 远程执行remote_entity，最终会执行remote.execute
             execution = remote.execute(
                entity,
                inputs=inputs,
                project=project,
                domain=domain,
                execution_name=run_level_params.name,
                options=options_from_run_params(run_level_params),
                type_hints=type_hints,
                overwrite_cache=run_level_params.overwrite_cache,
                interruptible=run_level_params.interruptible,
                envs=run_level_params.envvars,
                tags=run_level_params.tags,
                cluster_pool=run_level_params.cluster_pool,
                execution_cluster_label=run_level_params.execution_cluster_label,
            )

# 列举几个关键方法
class FlyteContextManager(object):

     @staticmethod
    def initialize():
        """
        Re-initializes the context and erases the entire context
        """
        # 设置默认id
        default_execution_id = _identifier.WorkflowExecutionIdentifier(project="local", domain="local", name="local")

        cfg = Config.auto()
        # 默认context工作目录.
        user_space_path = os.path.join(cfg.local_sandbox_path, "user_space")
        pathlib.Path(user_space_path).mkdir(parents=True, exist_ok=True)

        # 处理退出信号
        if threading.current_thread().name == threading.main_thread().name:
            signal.signal(signal.SIGINT, main_signal_handler)

        # 创建临时目录，并关联其provider，构建此栈帧的FlyteContext
        default_context = FlyteContext(file_access=default_local_file_access_provider)
        # 设置栈帧维护的参数
        default_user_space_params = ExecutionParameters(
            execution_id=WorkflowExecutionIdentifier.promote_from_model(default_execution_id),
            task_id=_identifier.Identifier(_identifier.ResourceType.TASK, "local", "local", "local", "local"),
            execution_date=datetime.now(timezone.utc),
            stats=mock_stats.MockStats(),
            logging=user_space_logger,
            tmp_dir=user_space_path,
            raw_output_prefix=default_context.file_access._raw_output_prefix,
            decks=[],
        )

        default_context = default_context.with_execution_state(
            default_context.new_execution_state().with_params(user_space_params=default_user_space_params)
        ).build()
        default_context.set_stackframe(s=FlyteContextManager.get_origin_stackframe())
        # 入栈
        flyte_context_Var.set([default_context])

    @staticmethod
    def current_context() -> FlyteContext:
        if not flyte_context_Var.get():
            # 默认创建一个default stack
            FlyteContextManager.initialize()
        # 返回栈顶context
        return flyte_context_Var.get()[-1]

    @staticmethod
    def get_origin_stackframe(limit=2) -> traceback.FrameSummary: 
        ss = traceback.extract_stack(limit=limit + 1) 
        if len(ss) > limit + 1:
            # 选择跳过的最后一个栈帧
            return ss[limit] 
        return ss[0]

    @staticmethod
    @contextmanager
    def with_context(b: FlyteContext.Builder) -> Generator[FlyteContext, None, None]:
        # b.build() 会创建一个当前的FlyteContext，并入栈，同时关联当前代码的stackframe(返回最近的3个调用栈frame)，跳过get_origin_stackframe调用函数
        ctx = FlyteContextManager.push_context(b.build(), FlyteContextManager.get_origin_stackframe(limit=3))
        l = FlyteContextManager.size()
        try:
            # 所有栈执行完后，退出上下文
            yield ctx
        finally:
            while FlyteContextManager.size() >= l:
                FlyteContextManager.pop_context()

if isinstance(entity, WorkflowBase):
    return self.execute_local_workflow(
        entity=entity,
        inputs=inputs,
        project=project,
        domain=domain,
        name=name,
        version=version,
        execution_name=execution_name,
        execution_name_prefix=execution_name_prefix,
        image_config=image_config,
        options=options,
        wait=wait,
        overwrite_cache=overwrite_cache,
        interruptible=interruptible,
        envs=envs,
        tags=tags,
        cluster_pool=cluster_pool,
        execution_cluster_label=execution_cluster_label,
        serialization_settings=serialization_settings,
    )

 def execute_local_workflow(
        self,
        entity: WorkflowBase,
        inputs: typing.Optional[typing.Dict[str, typing.Any]] = None,
        project: str = None,
        domain: str = None,
        name: str = None,
        version: str = None,
        execution_name: typing.Optional[str] = None,
        execution_name_prefix: typing.Optional[str] = None,
        image_config: typing.Optional[ImageConfig] = None,
        options: typing.Optional[Options] = None,
        wait: bool = False,
        overwrite_cache: typing.Optional[bool] = None,
        interruptible: typing.Optional[bool] = None,
        envs: typing.Optional[typing.Dict[str, str]] = None,
        tags: typing.Optional[typing.List[str]] = None,
        cluster_pool: typing.Optional[str] = None,
        execution_cluster_label: typing.Optional[str] = None,
        serialization_settings: typing.Optional[SerializationSettings] = None,
    ) -> FlyteWorkflowExecution:
        # 补全默认镜像
        if not image_config:
            image_config = ImageConfig.auto_default_image()
        # 如果没有序列号设置，补全
        ss = serialization_settings or SerializationSettings(
            image_config=image_config,
            project=project or self.default_project,
            domain=domain or self._default_domain,
            version=version,
        )
        if version is None and self.interactive_mode_enabled:
            version, _ = self._resolve_version(version, entity, ss)

        # input转换成字典
        resolved_identifiers = self._resolve_identifier_kwargs(entity, project, domain, name, version)
        resolved_identifiers_dict = asdict(resolved_identifiers)

        try:
            # Just fetch to see if it already exists
            # todo: Add logic to check that the fetched workflow is functionally equivalent.
            self.fetch_workflow(**resolved_identifiers_dict)
        except FlyteEntityNotExistException:
            logger.info("Registering workflow because it wasn't found in Flyte Admin.")
            # 如没有发现相关wf，重新注册
            self.register_workflow(
                entity,
                ss,
                version=version,
                options=options,
            )

        try:
            # 关联LP
            flyte_lp = self.fetch_launch_plan(**resolved_identifiers_dict)
            flyte_lp.python_interface = entity.python_interface
        except FlyteEntityNotExistException:
            logger.info("Try to register default launch plan because it wasn't found in Flyte Admin!")
            default_lp = LaunchPlan.get_default_launch_plan(self.context, entity)
            self.register_launch_plan(
                default_lp,
                project=resolved_identifiers.project,
                domain=resolved_identifiers.domain,
                version=version,
                options=options,
            )
            flyte_lp = self.fetch_launch_plan(**resolved_identifiers_dict)
        # 执行
        return self.execute(
            flyte_lp,
            inputs,
            project=project,
            domain=domain,
            execution_name=execution_name,
            execution_name_prefix=execution_name_prefix,
            wait=wait,
            options=options,
            type_hints=entity.python_interface.inputs,
            overwrite_cache=overwrite_cache,
            interruptible=interruptible,
            envs=envs,
            tags=tags,
            cluster_pool=cluster_pool,
            execution_cluster_label=execution_cluster_label,
        )

def _execute(self, ...) -> FlyteWorkflowExecution:
    # 根据当前文件切换上下文，获取文件上传地址
    self.file_access._get_upload_signed_url_fn = functools.partial(
        self.client.get_upload_signed_url,
        project=project or self.default_project,
        domain=domain or self.default_domain,
    )
    # 转换python input map 到 定义的 literal_inputs pb
    literal_inputs = literal_models.LiteralMap(literals=literal_map)

    # 发送执行请求
    exec_id = self.client.create_execution(
        project or self.default_project,
        domain or self.default_domain,
        execution_name,
        # 构建Spec和metadata
        ExecutionSpec(
            entity.id,
            ExecutionMetadata(
                ExecutionMetadata.ExecutionMode.MANUAL,
                "placeholder",  # Admin replaces this from oidc token if auth is enabled.
                0,
            ),
            overwrite_cache=overwrite_cache,
            interruptible=interruptible,
            notifications=notifications,
            disable_all=options.disable_notifications,
            labels=options.labels,
            annotations=options.annotations,
            raw_output_data_config=options.raw_output_data_config,
            auth_role=None,
            max_parallelism=options.max_parallelism,
            security_context=options.security_context,
            envs=common_models.Envs(envs) if envs else None,
            tags=tags,
            cluster_assignment=ClusterAssignment(cluster_pool=cluster_pool) if cluster_pool else None,
            execution_cluster_label=ExecutionClusterLabel(execution_cluster_label)
            if execution_cluster_label
            else None,
        ),
        literal_inputs,
    )

 exe = remote.execute(hello, inputs={"name": "Flyte"})
 # 并可以同步等待wf执行结果(1s检查一次)
 exe = exe.wait(poll_interval=1)