# Copyright (c) Microsoft. All rights reserved.

import asyncio
import inspect
import logging
from collections.abc import Callable, Sequence
from typing import Any, Literal, cast

from agent_framework import AgentResponse, Message, SupportsAgentRun
from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
from agent_framework._workflows._agent_utils import resolve_agent_id
from agent_framework._workflows._checkpoint import CheckpointStorage
from agent_framework._workflows._executor import Executor, handler
from agent_framework._workflows._message_utils import normalize_messages_input
from agent_framework._workflows._workflow import Workflow
from agent_framework._workflows._workflow_builder import WorkflowBuilder
from agent_framework._workflows._workflow_context import WorkflowContext
from typing_extensions import Never

from ._orchestration_request_info import AgentApprovalExecutor
from ._participant_output_config import (
    UNSET,
    _coalesce_output_from,  # pyright: ignore[reportPrivateUsage]
    _coerce_intermediate_output_from,  # pyright: ignore[reportPrivateUsage]
    _ParticipantIntermediateOutputSelection,  # pyright: ignore[reportPrivateUsage]
    _ParticipantOutputSpecifier,  # pyright: ignore[reportPrivateUsage]
    _resolve_participant_output_config,  # pyright: ignore[reportPrivateUsage]
)

logger = logging.getLogger(__name__)

"""Concurrent builder for agent-only fan-out/fan-in workflows.

This module provides a high-level, agent-focused API to quickly assemble a
parallel workflow with:
- a default dispatcher that broadcasts the input to all agent participants
- a default aggregator that combines all agent conversations and completes the workflow

Notes:
- Participants can be provided as SupportsAgentRun or Executor instances via `participants=[...]`.
- A custom aggregator can be provided as:
  - an Executor instance (it should handle list[AgentExecutorResponse],
    yield output), or
  - a callback function with signature:
        def cb(results: list[AgentExecutorResponse]) -> Any | None
        def cb(results: list[AgentExecutorResponse], ctx: WorkflowContext) -> Any | None
    The callback is wrapped in _CallbackAggregator.
    If the callback returns a non-None value, _CallbackAggregator yields that as output.
    If it returns None, the callback may have already yielded an output via ctx, so no further action is taken.
"""


class _DispatchToAllParticipants(Executor):
    """Broadcasts input to all downstream participants (via fan-out edges)."""

    @handler
    async def from_request(self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
        # No explicit target: edge routing delivers to all connected participants.
        await ctx.send_message(request)

    @handler
    async def from_str(self, prompt: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
        request = AgentExecutorRequest(messages=normalize_messages_input(prompt), should_respond=True)
        await ctx.send_message(request)

    @handler
    async def from_message(self, message: Message, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
        request = AgentExecutorRequest(messages=normalize_messages_input(message), should_respond=True)
        await ctx.send_message(request)

    @handler
    async def from_messages(
        self,
        messages: list[str | Message],
        ctx: WorkflowContext[AgentExecutorRequest],
    ) -> None:
        request = AgentExecutorRequest(messages=normalize_messages_input(messages), should_respond=True)
        await ctx.send_message(request)


class _AggregateAgentConversations(Executor):
    """Aggregates agent responses and completes with a single AgentResponse.

    Emits an `AgentResponse` whose `messages` are the final assistant message from each
    participant (one message per agent), in deterministic participant order matching
    the fan-in `sources` configuration. The user prompt is intentionally not included —
    that is part of the input, not the answer.

    For each participant the final assistant message is sourced from
    `r.agent_response.messages`, falling back to scanning `r.full_conversation` for
    pathological executors that did not populate the response.
    """

    @handler
    async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, AgentResponse]) -> None:
        if not results:
            logger.error("Concurrent aggregator received empty results list")
            raise ValueError("Aggregation failed: no results provided")

        def _is_role(msg: Any, role: str) -> bool:
            r = getattr(msg, "role", None)
            if r is None:
                return False
            r_str = str(r).lower() if isinstance(r, str) or hasattr(r, "__str__") else r
            role_str = str(role).lower()
            return r_str == role_str

        assistant_replies: list[Message] = []

        for r in results:
            resp_messages = list(r.agent_response.messages)

            logger.debug(
                f"Aggregating executor {getattr(r, 'executor_id', '<unknown>')}: "
                f"{len(resp_messages)} response msgs, {len(r.full_conversation)} conversation msgs"
            )

            # Pick the final assistant message from the response; fallback to conversation search
            final_assistant = next((m for m in reversed(resp_messages) if _is_role(m, "assistant")), None)
            if final_assistant is None:
                final_assistant = next((m for m in reversed(r.full_conversation) if _is_role(m, "assistant")), None)

            if final_assistant is not None:
                assistant_replies.append(final_assistant)
            else:
                logger.warning(
                    f"No assistant reply found for executor {getattr(r, 'executor_id', '<unknown>')}; skipping"
                )

        if not assistant_replies:
            logger.error(f"Aggregation failed: no assistant replies found across {len(results)} results")
            raise RuntimeError("Aggregation failed: no assistant replies found")

        await ctx.yield_output(AgentResponse(messages=assistant_replies))


class _CallbackAggregator(Executor):
    """Wraps a Python callback as an aggregator.

    Accepts either an async or sync callback with one of the signatures:
      - (results: list[AgentExecutorResponse]) -> Any | None
      - (results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> Any | None

    Notes:
    - Async callbacks are awaited directly.
    - Sync callbacks are executed via asyncio.to_thread to avoid blocking the event loop.
    - If the callback returns a non-None value, it is yielded as an output.
    """

    def __init__(self, callback: Callable[..., Any], id: str | None = None) -> None:
        derived_id = getattr(callback, "__name__", "") or ""
        if not derived_id or derived_id == "<lambda>":
            derived_id = f"{type(self).__name__}_unnamed"
        super().__init__(id or derived_id)
        self._callback = callback
        self._param_count = len(inspect.signature(callback).parameters)

    @handler
    async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, Any]) -> None:
        # Call according to provided signature, always non-blocking for sync callbacks
        if self._param_count >= 2:
            if inspect.iscoroutinefunction(self._callback):
                ret = await self._callback(results, ctx)
            else:
                ret = await asyncio.to_thread(self._callback, results, ctx)
        else:
            if inspect.iscoroutinefunction(self._callback):
                ret = await self._callback(results)
            else:
                ret = await asyncio.to_thread(self._callback, results)

        # If the callback returned a value, finalize the workflow with it
        if ret is not None:
            await ctx.yield_output(ret)


class ConcurrentBuilder:
    r"""High-level builder for concurrent agent workflows.

    - `participants=[...]` accepts a list of SupportsAgentRun (recommended) or Executor.
    - `build()` wires: dispatcher -> fan-out -> participants -> fan-in -> aggregator.
    - `with_aggregator(...)` overrides the default aggregator with an Executor or callback.

    Usage:

    .. code-block:: python

        from agent_framework_orchestrations import ConcurrentBuilder

        # Minimal: use default aggregator (yields one AgentResponse with one assistant
        # message per participant)
        workflow = ConcurrentBuilder(participants=[agent1, agent2, agent3]).build()


        # Custom aggregator via callback (sync or async). The callback receives
        # list[AgentExecutorResponse] and its return value becomes the workflow's output.
        def summarize(results: list[AgentExecutorResponse]) -> str:
            return " | ".join(r.agent_response.messages[-1].text for r in results)


        workflow = ConcurrentBuilder(participants=[agent1, agent2, agent3]).with_aggregator(summarize).build()


        # Enable checkpoint persistence so runs can resume
        workflow = ConcurrentBuilder(participants=[agent1, agent2, agent3], checkpoint_storage=storage).build()

        # Enable request info before aggregation
        workflow = ConcurrentBuilder(participants=[agent1, agent2]).with_request_info().build()
    """

    def __init__(
        self,
        *,
        participants: Sequence[SupportsAgentRun | Executor],
        checkpoint_storage: CheckpointStorage | None = None,
        output_from: Sequence[_ParticipantOutputSpecifier] | Literal["all"] | None = cast(Any, UNSET),
        intermediate_output_from: _ParticipantIntermediateOutputSelection = None,
    ) -> None:
        """Initialize the ConcurrentBuilder.

        Args:
            participants: Sequence of agent or executor instances to run in parallel.
            checkpoint_storage: Optional checkpoint storage for enabling workflow state persistence.
            output_from: Optional participant names or instances whose ``yield_output`` calls
                surface as workflow ``output`` events alongside the aggregator. Pass ``"all"`` to select every
                participant.
            intermediate_output_from: Optional participant names or instances whose ``yield_output`` calls
                surface as workflow ``intermediate`` events. Pass ``"all_other"`` to select every participant
                not selected by ``output_from``. Unlisted participant outputs are hidden.
        """
        self._participants: list[SupportsAgentRun | Executor] = []
        self._aggregator: Executor | None = None
        self._checkpoint_storage: CheckpointStorage | None = checkpoint_storage
        self._request_info_enabled: bool = False
        self._request_info_filter: set[str] | None = None
        self._output_from = _coalesce_output_from(output_from=output_from)
        self._intermediate_output_from = _coerce_intermediate_output_from(intermediate_output_from)

        self._set_participants(participants)

    def _set_participants(self, participants: Sequence[SupportsAgentRun | Executor]) -> None:
        """Set participants (internal)."""
        if self._participants:
            raise ValueError("participants already set.")

        if not participants:
            raise ValueError("participants cannot be empty")

        # Defensive duplicate detection
        seen_agent_ids: set[int] = set()
        seen_executor_ids: set[str] = set()
        for p in participants:
            if isinstance(p, Executor):
                if p.id in seen_executor_ids:
                    raise ValueError(f"Duplicate executor participant detected: id '{p.id}'")
                seen_executor_ids.add(p.id)
            elif isinstance(p, SupportsAgentRun):
                pid = id(p)
                if pid in seen_agent_ids:
                    raise ValueError("Duplicate agent participant detected (same agent instance provided twice)")
                seen_agent_ids.add(pid)
            else:
                raise TypeError(f"participants must be SupportsAgentRun or Executor instances; got {type(p).__name__}")

        self._participants = list(participants)

    def with_aggregator(
        self,
        aggregator: Executor
        | Callable[[list[AgentExecutorResponse]], Any]
        | Callable[[list[AgentExecutorResponse], WorkflowContext[Never, Any]], Any],
    ) -> "ConcurrentBuilder":
        r"""Override the default aggregator with an executor or a callback.

        - Executor: must handle `list[AgentExecutorResponse]` and yield output using `ctx.yield_output(...)`
        - Callback: sync or async callable with one of the signatures:
          `(results: list[AgentExecutorResponse]) -> Any | None` or
          `(results: list[AgentExecutorResponse], ctx: WorkflowContext) -> Any | None`.
          If the callback returns a non-None value, it becomes the workflow's output.

        Args:
            aggregator: Executor instance, or callback function

        Example:

        .. code-block:: python
            # Executor-based aggregator
            class CustomAggregator(Executor):
                @handler
                async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext) -> None:
                    await ctx.yield_output(" | ".join(r.agent_response.messages[-1].text for r in results))


            wf = ConcurrentBuilder(participants=[a1, a2, a3]).with_aggregator(CustomAggregator()).build()


            # Callback-based aggregator (string result)
            async def summarize(results: list[AgentExecutorResponse]) -> str:
                return " | ".join(r.agent_response.messages[-1].text for r in results)


            wf = ConcurrentBuilder(participants=[a1, a2, a3]).with_aggregator(summarize).build()


            # Callback-based aggregator (yield result)
            async def summarize(results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, str]) -> None:
                await ctx.yield_output(" | ".join(r.agent_response.messages[-1].text for r in results))


            wf = ConcurrentBuilder(participants=[a1, a2, a3]).with_aggregator(summarize).build()
        """
        if self._aggregator is not None:
            raise ValueError("with_aggregator() has already been called on this builder instance.")

        if isinstance(aggregator, Executor):
            self._aggregator = aggregator
        elif callable(aggregator):
            self._aggregator = _CallbackAggregator(aggregator)
        else:
            raise TypeError("aggregator must be an Executor or a callable")

        return self

    def with_request_info(
        self,
        *,
        agents: Sequence[str | SupportsAgentRun] | None = None,
    ) -> "ConcurrentBuilder":
        """Enable request info after agent participant responses.

        This enables human-in-the-loop (HIL) scenarios for the concurrent orchestration.
        When enabled, the workflow pauses after each agent participant runs, emitting
        a request_info event (type='request_info') that allows the caller to review the conversation and optionally
        inject guidance for the agent participant to iterate. The caller provides input via
        the standard response_handler/request_info pattern.

        Simulated flow with HIL:
        Input -> [Agent Participant <-> Request Info] -> [Agent Participant <-> Request Info] -> ...

        Note: This is only available for agent participants. Executor participants can incorporate
        request info handling in their own implementation if desired.

        Args:
            agents: Optional list of agents names or agent factories to enable request info for.
                    If None, enables HIL for all agent participants.

        Returns:
            Self for fluent chaining
        """
        from ._orchestration_request_info import resolve_request_info_filter

        self._request_info_enabled = True
        self._request_info_filter = resolve_request_info_filter(list(agents) if agents else None)

        return self

    def _resolve_participants(self) -> list[Executor]:
        """Resolve participant instances into Executor objects."""
        if not self._participants:
            raise ValueError("No participants provided. Pass participants to the constructor.")

        participants: list[Executor | SupportsAgentRun] = self._participants

        executors: list[Executor] = []
        for p in participants:
            if isinstance(p, Executor):
                executors.append(p)
            elif isinstance(p, SupportsAgentRun):
                if self._request_info_enabled and (
                    not self._request_info_filter or resolve_agent_id(p) in self._request_info_filter
                ):
                    # Handle request info enabled agents
                    executors.append(AgentApprovalExecutor(p))
                else:
                    executors.append(AgentExecutor(p))
            else:
                raise TypeError(f"Participants must be SupportsAgentRun or Executor instances. Got {type(p).__name__}.")

        return executors

    def build(self) -> Workflow:
        r"""Build and validate the concurrent workflow.

        Wiring pattern:
        - Dispatcher (internal) fans out the input to all `participants`
        - Fan-in collects `AgentExecutorResponse` objects from all participants
        - If request info is enabled, the orchestration emits a request info event with outputs from all participants
            before sending the outputs to the aggregator
        - Aggregator yields output and the workflow becomes idle. The output is either:
          - AgentResponse (default aggregator: one assistant message per participant)
          - custom payload from the provided aggregator

        Returns:
            Workflow: a ready-to-run workflow instance

        Raises:
            ValueError: if no participants were defined

        Example:

        .. code-block:: python

            workflow = ConcurrentBuilder(participants=[agent1, agent2]).build()
        """
        # Internal nodes
        dispatcher = _DispatchToAllParticipants(id="dispatcher")
        aggregator = self._aggregator if self._aggregator is not None else _AggregateAgentConversations(id="aggregator")

        # Resolve participants and participant factories to executors
        participants: list[Executor] = self._resolve_participants()

        # Default: only the aggregator is terminal; participant outputs are hidden
        # unless explicitly designated as terminal or intermediate.
        designated, intermediate_designated = _resolve_participant_output_config(
            participants=participants,
            output_from=self._output_from,
            intermediate_output_from=self._intermediate_output_from,
            extra_output_executors=[aggregator],
        )
        builder = WorkflowBuilder(
            start_executor=dispatcher,
            checkpoint_storage=self._checkpoint_storage,
            output_from=designated,
            intermediate_output_from=intermediate_designated,
        )
        # Fan-out for parallel execution
        builder.add_fan_out_edges(dispatcher, participants)
        # Direct fan-in to aggregator
        builder.add_fan_in_edges(participants, aggregator)

        return builder.build()
