import asyncio
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterator, List, Optional, Union
from uuid import uuid4

from agno.models.metrics import Metrics
from agno.run.agent import RunOutputEvent
from agno.run.team import TeamRunOutputEvent
from agno.run.workflow import (
    ParallelExecutionCompletedEvent,
    ParallelExecutionStartedEvent,
    WorkflowRunOutput,
    WorkflowRunOutputEvent,
)
from agno.utils.log import log_debug, logger
from agno.workflow.condition import Condition
from agno.workflow.step import Step
from agno.workflow.types import StepInput, StepOutput, StepType

WorkflowSteps = List[
    Union[
        Callable[
            [StepInput], Union[StepOutput, Awaitable[StepOutput], Iterator[StepOutput], AsyncIterator[StepOutput]]
        ],
        Step,
        "Steps",  # type: ignore # noqa: F821
        "Loop",  # type: ignore # noqa: F821
        "Parallel",  # type: ignore # noqa: F821
        "Condition",  # type: ignore # noqa: F821
        "Router",  # type: ignore # noqa: F821
    ]
]


@dataclass
class Parallel:
    """A list of steps that execute in parallel"""

    steps: WorkflowSteps

    name: Optional[str] = None
    description: Optional[str] = None

    def __init__(
        self,
        *steps: WorkflowSteps,
        name: Optional[str] = None,
        description: Optional[str] = None,
    ):
        self.steps = list(steps)
        self.name = name
        self.description = description

    def _prepare_steps(self):
        """Prepare the steps for execution - mirrors workflow logic"""
        from agno.agent.agent import Agent
        from agno.team.team import Team
        from agno.workflow.loop import Loop
        from agno.workflow.router import Router
        from agno.workflow.step import Step
        from agno.workflow.steps import Steps

        prepared_steps: WorkflowSteps = []
        for step in self.steps:
            if callable(step) and hasattr(step, "__name__"):
                prepared_steps.append(Step(name=step.__name__, description="User-defined callable step", executor=step))
            elif isinstance(step, Agent):
                prepared_steps.append(Step(name=step.name, description=step.description, agent=step))
            elif isinstance(step, Team):
                prepared_steps.append(Step(name=step.name, description=step.description, team=step))
            elif isinstance(step, (Step, Steps, Loop, Parallel, Condition, Router)):
                prepared_steps.append(step)
            else:
                raise ValueError(f"Invalid step type: {type(step).__name__}")

        self.steps = prepared_steps

    def _aggregate_results(self, step_outputs: List[StepOutput]) -> StepOutput:
        """Aggregate multiple step outputs into a single StepOutput"""
        if not step_outputs:
            return StepOutput(
                step_name=self.name or "Parallel",
                step_id=str(uuid4()),
                step_type="Parallel",
                content="No parallel steps executed",
                steps=[],
            )

        if len(step_outputs) == 1:
            # Single result, but still create a Parallel container
            single_result = step_outputs[0]
            aggregated_metrics = self._extract_metrics_from_response(step_outputs)

            return StepOutput(
                step_name=self.name or "Parallel",
                step_id=str(uuid4()),
                step_type=StepType.PARALLEL,
                content=f"Parallel {self.name or 'execution'} completed with 1 result",
                executor_name=self.name or "Parallel",
                images=single_result.images,
                videos=single_result.videos,
                audio=single_result.audio,
                metrics=aggregated_metrics,
                success=single_result.success,
                error=single_result.error,
                stop=single_result.stop,
                steps=step_outputs,  # This is the key addition
            )

        early_termination_requested = any(output.stop for output in step_outputs if hasattr(output, "stop"))

        # Multiple results - aggregate them
        aggregated_content = f"Parallel {self.name or 'execution'} completed with {len(step_outputs)} results"

        # Combine all media from parallel steps
        all_images = []
        all_videos = []
        all_audio = []
        has_any_failure = False

        for result in step_outputs:
            all_images.extend(result.images or [])
            all_videos.extend(result.videos or [])
            all_audio.extend(result.audio or [])
            if result.success is False:
                has_any_failure = True

        # Extract metrics using the dedicated method
        aggregated_metrics = self._extract_metrics_from_response(step_outputs)

        return StepOutput(
            step_name=self.name or "Parallel",
            step_id=str(uuid4()),
            step_type=StepType.PARALLEL,
            executor_type="parallel",
            executor_name=self.name or "Parallel",
            content=aggregated_content,
            images=all_images if all_images else None,
            videos=all_videos if all_videos else None,
            audio=all_audio if all_audio else None,
            success=not has_any_failure,
            stop=early_termination_requested,
            metrics=aggregated_metrics,
            steps=step_outputs,
        )

    def _extract_metrics_from_response(self, step_outputs: List[StepOutput]) -> Optional[Metrics]:
        """Extract and aggregate metrics from parallel step outputs"""
        if not step_outputs:
            return None

        # Aggregate metrics from all parallel step outputs
        total_metrics = Metrics()

        for result in step_outputs:
            if result.metrics:
                total_metrics = total_metrics + result.metrics

        # If no metrics were found, return None
        if (
            total_metrics.input_tokens == 0
            and total_metrics.output_tokens == 0
            and total_metrics.total_tokens == 0
            and total_metrics.duration is None
        ):
            return None

        return total_metrics

    def _build_aggregated_content(self, step_outputs: List[StepOutput]) -> str:
        """Build aggregated content from multiple step outputs"""
        aggregated = "## Parallel Execution Results\n\n"

        for i, output in enumerate(step_outputs):
            step_name = output.step_name or f"Step {i + 1}"
            content = output.content or ""

            # Add status indicator
            if output.success is False:
                status_icon = "❌ FAILURE:"
            else:
                status_icon = "✅ SUCCESS:"

            aggregated += f"### {status_icon} {step_name}\n"
            if content and str(content).strip():
                aggregated += f"{content}\n\n"
            else:
                aggregated += "*(No content)*\n\n"

        return aggregated.strip()

    def execute(
        self,
        step_input: StepInput,
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        workflow_run_response: Optional[WorkflowRunOutput] = None,
        store_executor_outputs: bool = True,
        session_state: Optional[Dict[str, Any]] = None,
    ) -> StepOutput:
        """Execute all steps in parallel and return aggregated result"""
        # Use workflow logger for parallel orchestration
        log_debug(f"Parallel Start: {self.name} ({len(self.steps)} steps)", center=True, symbol="=")

        self._prepare_steps()

        def execute_step_with_index(step_with_index):
            """Execute a single step and preserve its original index"""
            idx, step = step_with_index
            try:
                step_result = step.execute(
                    step_input,
                    session_id=session_id,
                    user_id=user_id,
                    workflow_run_response=workflow_run_response,
                    store_executor_outputs=store_executor_outputs,
                    session_state=session_state,
                )  # type: ignore[union-attr]
                return idx, step_result
            except Exception as exc:
                parallel_step_name = getattr(step, "name", f"step_{idx}")
                logger.error(f"Parallel step {parallel_step_name} failed: {exc}")
                return (
                    idx,
                    StepOutput(
                        step_name=parallel_step_name,
                        content=f"Step {parallel_step_name} failed: {str(exc)}",
                        success=False,
                        error=str(exc),
                    ),
                )

        # Use index to preserve order
        indexed_steps = list(enumerate(self.steps))

        with ThreadPoolExecutor(max_workers=len(self.steps)) as executor:
            # Submit all tasks with their original indices
            future_to_index = {
                executor.submit(execute_step_with_index, indexed_step): indexed_step[0]
                for indexed_step in indexed_steps
            }

            # Collect results
            results_with_indices = []
            for future in as_completed(future_to_index):
                try:
                    index, result = future.result()
                    results_with_indices.append((index, result))
                    step_name = getattr(self.steps[index], "name", f"step_{index}")
                    log_debug(f"Parallel step {step_name} completed")
                except Exception as e:
                    index = future_to_index[future]
                    step_name = getattr(self.steps[index], "name", f"step_{index}")
                    logger.error(f"Parallel step {step_name} failed: {e}")
                    results_with_indices.append(
                        (
                            index,
                            StepOutput(
                                step_name=step_name,
                                content=f"Step {step_name} failed: {str(e)}",
                                success=False,
                                error=str(e),
                            ),
                        )
                    )

        # Sort by original index to preserve order
        results_with_indices.sort(key=lambda x: x[0])
        results = [result for _, result in results_with_indices]

        # Flatten results - handle steps that return List[StepOutput] (like Condition/Loop)
        flattened_results: List[StepOutput] = []
        for result in results:
            if isinstance(result, list):
                flattened_results.extend(result)
            else:
                flattened_results.append(result)

        # Aggregate all results into a single StepOutput
        aggregated_result = self._aggregate_results(flattened_results)

        # Use workflow logger for parallel completion
        log_debug(f"Parallel End: {self.name} ({len(self.steps)} steps)", center=True, symbol="=")

        return aggregated_result

    def execute_stream(
        self,
        step_input: StepInput,
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        stream_intermediate_steps: bool = False,
        workflow_run_response: Optional[WorkflowRunOutput] = None,
        step_index: Optional[Union[int, tuple]] = None,
        store_executor_outputs: bool = True,
        session_state: Optional[Dict[str, Any]] = None,
        parent_step_id: Optional[str] = None,
    ) -> Iterator[Union[WorkflowRunOutputEvent, StepOutput]]:
        """Execute all steps in parallel with streaming support"""
        log_debug(f"Parallel Start: {self.name} ({len(self.steps)} steps)", center=True, symbol="=")

        parallel_step_id = str(uuid4())

        self._prepare_steps()

        if stream_intermediate_steps and workflow_run_response:
            # Yield parallel step started event
            yield ParallelExecutionStartedEvent(
                run_id=workflow_run_response.run_id or "",
                workflow_name=workflow_run_response.workflow_name or "",
                workflow_id=workflow_run_response.workflow_id or "",
                session_id=workflow_run_response.session_id or "",
                step_name=self.name,
                step_index=step_index,
                parallel_step_count=len(self.steps),
                step_id=parallel_step_id,
                parent_step_id=parent_step_id,
            )

        def execute_step_stream_with_index(step_with_index):
            """Execute a single step with streaming and preserve its original index"""
            idx, step = step_with_index
            try:
                step_events = []

                # If step_index is None or integer (main step): create (step_index, sub_index)
                # If step_index is tuple (child step): all parallel sub-steps get same index
                if step_index is None or isinstance(step_index, int):
                    # Parallel is a main step - sub-steps get sequential numbers: 1.1, 1.2, 1.3
                    sub_step_index = (step_index if step_index is not None else 0, idx)
                else:
                    # Parallel is a child step - all sub-steps get the same parent number: 1.1, 1.1, 1.1
                    sub_step_index = step_index

                # All workflow step types have execute_stream() method
                for event in step.execute_stream(  # type: ignore[union-attr]
                    step_input,
                    session_id=session_id,
                    user_id=user_id,
                    stream_intermediate_steps=stream_intermediate_steps,
                    workflow_run_response=workflow_run_response,
                    step_index=sub_step_index,
                    store_executor_outputs=store_executor_outputs,
                    session_state=session_state,
                    parent_step_id=parallel_step_id,
                ):
                    step_events.append(event)
                return idx, step_events
            except Exception as exc:
                parallel_step_name = getattr(step, "name", f"step_{idx}")
                logger.error(f"Parallel step {parallel_step_name} streaming failed: {exc}")
                return (
                    idx,
                    [
                        StepOutput(
                            step_name=parallel_step_name,
                            content=f"Step {parallel_step_name} failed: {str(exc)}",
                            success=False,
                            error=str(exc),
                        )
                    ],
                )

        # Use index to preserve order
        indexed_steps = list(enumerate(self.steps))
        all_events_with_indices = []
        step_results = []

        with ThreadPoolExecutor(max_workers=len(self.steps)) as executor:
            # Submit all tasks with their original indices
            future_to_index = {
                executor.submit(execute_step_stream_with_index, indexed_step): indexed_step[0]
                for indexed_step in indexed_steps
            }

            # Collect results as they complete
            for future in as_completed(future_to_index):
                try:
                    index, events = future.result()
                    all_events_with_indices.append((index, events))

                    # Extract StepOutput from events for the final result
                    step_outputs = [event for event in events if isinstance(event, StepOutput)]
                    if step_outputs:
                        step_results.extend(step_outputs)

                    step_name = getattr(self.steps[index], "name", f"step_{index}")
                    log_debug(f"Parallel step {step_name} streaming completed")
                except Exception as e:
                    index = future_to_index[future]
                    step_name = getattr(self.steps[index], "name", f"step_{index}")
                    logger.error(f"Parallel step {step_name} streaming failed: {e}")
                    error_event = StepOutput(
                        step_name=step_name,
                        content=f"Step {step_name} failed: {str(e)}",
                        success=False,
                        error=str(e),
                    )
                    all_events_with_indices.append((index, [error_event]))
                    step_results.append(error_event)

        # Sort events by original index to preserve order
        all_events_with_indices.sort(key=lambda x: x[0])

        # Yield all collected streaming events in order (but not final StepOutputs)
        for _, events in all_events_with_indices:
            for event in events:
                # Only yield non-StepOutput events during streaming to avoid duplication
                if not isinstance(event, StepOutput):
                    yield event

        # Flatten step_results - handle steps that return List[StepOutput] (like Condition/Loop)
        flattened_step_results: List[StepOutput] = []
        for result in step_results:
            if isinstance(result, list):
                flattened_step_results.extend(result)
            else:
                flattened_step_results.append(result)

        # Create aggregated result from all step outputs
        aggregated_result = self._aggregate_results(flattened_step_results)

        # Yield the final aggregated StepOutput
        yield aggregated_result

        log_debug(f"Parallel End: {self.name} ({len(self.steps)} steps)", center=True, symbol="=")

        if stream_intermediate_steps and workflow_run_response:
            # Yield parallel step completed event
            yield ParallelExecutionCompletedEvent(
                run_id=workflow_run_response.run_id or "",
                workflow_name=workflow_run_response.workflow_name or "",
                workflow_id=workflow_run_response.workflow_id or "",
                session_id=workflow_run_response.session_id or "",
                step_name=self.name,
                step_index=step_index,
                parallel_step_count=len(self.steps),
                step_results=[aggregated_result],  # Now single aggregated result
                step_id=parallel_step_id,
                parent_step_id=parent_step_id,
            )

    async def aexecute(
        self,
        step_input: StepInput,
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        workflow_run_response: Optional[WorkflowRunOutput] = None,
        store_executor_outputs: bool = True,
        session_state: Optional[Dict[str, Any]] = None,
    ) -> StepOutput:
        """Execute all steps in parallel using asyncio and return aggregated result"""
        # Use workflow logger for async parallel orchestration
        log_debug(f"Parallel Start: {self.name} ({len(self.steps)} steps)", center=True, symbol="=")

        self._prepare_steps()

        async def execute_step_async_with_index(step_with_index):
            """Execute a single step asynchronously and preserve its original index"""
            idx, step = step_with_index
            try:
                inner_step_result = await step.aexecute(
                    step_input,
                    session_id=session_id,
                    user_id=user_id,
                    workflow_run_response=workflow_run_response,
                    store_executor_outputs=store_executor_outputs,
                    session_state=session_state,
                )  # type: ignore[union-attr]
                return idx, inner_step_result
            except Exception as exc:
                parallel_step_name = getattr(step, "name", f"step_{idx}")
                logger.error(f"Parallel step {parallel_step_name} failed: {exc}")
                return (
                    idx,
                    StepOutput(
                        step_name=parallel_step_name,
                        content=f"Step {parallel_step_name} failed: {str(exc)}",
                        success=False,
                        error=str(exc),
                    ),
                )

        # Use index to preserve order
        indexed_steps = list(enumerate(self.steps))

        # Create tasks for all steps with their indices
        tasks = [execute_step_async_with_index(indexed_step) for indexed_step in indexed_steps]

        # Execute all tasks concurrently
        results_with_indices = await asyncio.gather(*tasks, return_exceptions=True)

        # Process results and handle exceptions, preserving order
        processed_results_with_indices = []
        for i, result in enumerate(results_with_indices):
            if isinstance(result, Exception):
                step_name = getattr(self.steps[i], "name", f"step_{i}")
                logger.error(f"Parallel step {step_name} failed: {result}")
                processed_results_with_indices.append(
                    (
                        i,
                        StepOutput(
                            step_name=step_name,
                            content=f"Step {step_name} failed: {str(result)}",
                            success=False,
                            error=str(result),
                        ),
                    )
                )
            else:
                index, step_result = result  # type: ignore[misc]
                processed_results_with_indices.append((index, step_result))
                step_name = getattr(self.steps[index], "name", f"step_{index}")
                log_debug(f"Parallel step {step_name} completed")

        # Sort by original index to preserve order
        processed_results_with_indices.sort(key=lambda x: x[0])
        results = [result for _, result in processed_results_with_indices]

        # Flatten results - handle steps that return List[StepOutput] (like Condition/Loop)
        flattened_results: List[StepOutput] = []
        for result in results:
            if isinstance(result, list):
                flattened_results.extend(result)
            else:
                flattened_results.append(result)

        # Aggregate all results into a single StepOutput
        aggregated_result = self._aggregate_results(flattened_results)

        # Use workflow logger for async parallel completion
        log_debug(f"Parallel End: {self.name} ({len(self.steps)} steps)", center=True, symbol="=")

        return aggregated_result

    async def aexecute_stream(
        self,
        step_input: StepInput,
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        stream_intermediate_steps: bool = False,
        workflow_run_response: Optional[WorkflowRunOutput] = None,
        step_index: Optional[Union[int, tuple]] = None,
        store_executor_outputs: bool = True,
        session_state: Optional[Dict[str, Any]] = None,
        parent_step_id: Optional[str] = None,
    ) -> AsyncIterator[Union[WorkflowRunOutputEvent, TeamRunOutputEvent, RunOutputEvent, StepOutput]]:
        """Execute all steps in parallel with async streaming support"""
        log_debug(f"Parallel Start: {self.name} ({len(self.steps)} steps)", center=True, symbol="=")

        parallel_step_id = str(uuid4())

        self._prepare_steps()

        if stream_intermediate_steps and workflow_run_response:
            # Yield parallel step started event
            yield ParallelExecutionStartedEvent(
                run_id=workflow_run_response.run_id or "",
                workflow_name=workflow_run_response.workflow_name or "",
                workflow_id=workflow_run_response.workflow_id or "",
                session_id=workflow_run_response.session_id or "",
                step_name=self.name,
                step_index=step_index,
                parallel_step_count=len(self.steps),
                step_id=parallel_step_id,
                parent_step_id=parent_step_id,
            )

        async def execute_step_stream_async_with_index(step_with_index):
            """Execute a single step with async streaming and preserve its original index"""
            idx, step = step_with_index
            try:
                step_events = []

                # If step_index is None or integer (main step): create (step_index, sub_index)
                # If step_index is tuple (child step): all parallel sub-steps get same index
                if step_index is None or isinstance(step_index, int):
                    # Parallel is a main step - sub-steps get sequential numbers: 1.1, 1.2, 1.3
                    sub_step_index = (step_index if step_index is not None else 0, idx)
                else:
                    # Parallel is a child step - all sub-steps get the same parent number: 1.1, 1.1, 1.1
                    sub_step_index = step_index

                # All workflow step types have aexecute_stream() method
                async for event in step.aexecute_stream(
                    step_input,
                    session_id=session_id,
                    user_id=user_id,
                    stream_intermediate_steps=stream_intermediate_steps,
                    workflow_run_response=workflow_run_response,
                    step_index=sub_step_index,
                    store_executor_outputs=store_executor_outputs,
                    session_state=session_state,
                    parent_step_id=parallel_step_id,
                ):  # type: ignore[union-attr]
                    step_events.append(event)
                return idx, step_events
            except Exception as e:
                parallel_step_name = getattr(step, "name", f"step_{idx}")
                logger.error(f"Parallel step {parallel_step_name} async streaming failed: {e}")
                return (
                    idx,
                    [
                        StepOutput(
                            step_name=parallel_step_name,
                            content=f"Step {parallel_step_name} failed: {str(e)}",
                            success=False,
                            error=str(e),
                        )
                    ],
                )

        # Use index to preserve order
        indexed_steps = list(enumerate(self.steps))
        all_events_with_indices = []
        step_results = []

        # Create tasks for all steps with their indices
        tasks = [execute_step_stream_async_with_index(indexed_step) for indexed_step in indexed_steps]

        # Execute all tasks concurrently
        results_with_indices = await asyncio.gather(*tasks, return_exceptions=True)

        # Process results and handle exceptions, preserving order
        for i, result in enumerate(results_with_indices):
            if isinstance(result, Exception):
                step_name = getattr(self.steps[i], "name", f"step_{i}")
                logger.error(f"Parallel step {step_name} async streaming failed: {result}")
                error_event = StepOutput(
                    step_name=step_name,
                    content=f"Step {step_name} failed: {str(result)}",
                    success=False,
                    error=str(result),
                )
                all_events_with_indices.append((i, [error_event]))
                step_results.append(error_event)
            else:
                index, events = result  # type: ignore[misc]
                all_events_with_indices.append((index, events))

                # Extract StepOutput from events for the final result
                step_outputs = [event for event in events if isinstance(event, StepOutput)]
                if step_outputs:
                    step_results.extend(step_outputs)

                step_name = getattr(self.steps[index], "name", f"step_{index}")
                log_debug(f"Parallel step {step_name} async streaming completed")

        # Sort events by original index to preserve order
        all_events_with_indices.sort(key=lambda x: x[0])

        # Yield all collected streaming events in order (but not final StepOutputs)
        for _, events in all_events_with_indices:
            for event in events:
                # Only yield non-StepOutput events during streaming to avoid duplication
                if not isinstance(event, StepOutput):
                    yield event

        # Flatten step_results - handle steps that return List[StepOutput] (like Condition/Loop)
        flattened_step_results: List[StepOutput] = []
        for result in step_results:
            if isinstance(result, list):
                flattened_step_results.extend(result)
            else:
                flattened_step_results.append(result)

        # Create aggregated result from all step outputs
        aggregated_result = self._aggregate_results(flattened_step_results)

        # Yield the final aggregated StepOutput
        yield aggregated_result

        log_debug(f"Parallel End: {self.name} ({len(self.steps)} steps)", center=True, symbol="=")

        if stream_intermediate_steps and workflow_run_response:
            # Yield parallel step completed event
            yield ParallelExecutionCompletedEvent(
                run_id=workflow_run_response.run_id or "",
                workflow_name=workflow_run_response.workflow_name or "",
                workflow_id=workflow_run_response.workflow_id or "",
                session_id=workflow_run_response.session_id or "",
                step_name=self.name,
                step_index=step_index,
                parallel_step_count=len(self.steps),
                step_results=[aggregated_result],  # Now single aggregated result
                step_id=parallel_step_id,
                parent_step_id=parent_step_id,
            )
