"""DynamicPipeline - 动态 Pipeline 组合

支持运行时根据条件选择子流程、嵌套 Pipeline、循环 Pipeline。
"""

import logging
from typing import Any

from agentkit.orchestrator.pipeline_engine import PipelineEngine
from agentkit.orchestrator.pipeline_schema import Pipeline, PipelineResult, StageStatus

logger = logging.getLogger(__name__)


class DynamicPipeline:
    """动态 Pipeline 组合器"""

    def __init__(self, engine: PipelineEngine, loader: Any = None):
        self._engine = engine
        self._loader = loader

    async def execute_conditional(
        self,
        pipelines: dict[str, Pipeline],
        condition_key: str,
        context: dict[str, Any] | None = None,
    ) -> PipelineResult:
        """根据条件选择子 Pipeline 执行"""
        context = context or {}
        condition_value = context.get(condition_key)

        if condition_value not in pipelines:
            return PipelineResult(
                pipeline_name=f"conditional_{condition_key}",
                status=StageStatus.FAILED,
                error_message=f"No pipeline for condition '{condition_key}={condition_value}'",
            )

        selected = pipelines[condition_value]
        logger.info(f"DynamicPipeline selected '{selected.name}' for {condition_key}={condition_value}")
        return await self._engine.execute(selected, context)

    async def execute_nested(
        self,
        parent: Pipeline,
        sub_pipeline_map: dict[str, Pipeline],
        context: dict[str, Any] | None = None,
    ) -> PipelineResult:
        """执行嵌套 Pipeline"""
        # 先执行父 Pipeline
        parent_result = await self._engine.execute(parent, context)

        # 根据父 Pipeline 结果选择子 Pipeline
        for stage_name, stage_result in parent_result.stage_results.items():
            if hasattr(stage_result, 'output_data') and stage_result.output_data:
                sub_pipeline_name = stage_result.output_data.get("sub_pipeline")
                if sub_pipeline_name and sub_pipeline_name in sub_pipeline_map:
                    sub = sub_pipeline_map[sub_pipeline_name]
                    sub_result = await self._engine.execute(sub, parent_result.variables)
                    parent_result.variables.update(sub_result.variables)

        return parent_result

    async def execute_loop(
        self,
        pipeline: Pipeline,
        max_iterations: int = 5,
        exit_condition: str = "done",
        context: dict[str, Any] | None = None,
    ) -> PipelineResult:
        """循环执行 Pipeline 直到条件满足"""
        current_context = context or {}
        last_result = None

        for i in range(max_iterations):
            logger.info(f"DynamicPipeline loop iteration {i + 1}/{max_iterations}")
            result = await self._engine.execute(pipeline, current_context)
            last_result = result

            # 检查退出条件
            if exit_condition in result.variables and result.variables[exit_condition]:
                logger.info(f"DynamicPipeline loop exited at iteration {i + 1}")
                break

            # 将结果作为下一轮的输入
            current_context.update(result.variables)

        return last_result or PipelineResult(
            pipeline_name=pipeline.name,
            status=StageStatus.FAILED,
            error_message="Loop completed without meeting exit condition",
        )
