"""Builder for preparing flow execution contexts.

This module isolates the complexity of setting up the execution environment,
variable initialization, and resource allocation (like MCP session pools).
It creates resources but transfers ownership to the executor.
"""

from typing import Any

from agent_flows.core.runtime import MCPSessionPool
from agent_flows.models.config import AgentFlowsConfig
from agent_flows.models.execution import ExecutionContext, TraceEntry
from agent_flows.models.flow import FlowConfig
from agent_flows.services.realtimex import VariableInitializer
from agent_flows.services.streaming import StreamingHandler


class ExecutionContextBuilder:
    """Builder for constructing ExecutionContext objects.

    Responsibilities:
        - Initialize flow variables
        - Create MCP session pools (lifecycle managed by caller)
        - Setup trace logging
        - Construct ExecutionContext with all dependencies
    """

    def __init__(
        self,
        config: AgentFlowsConfig,
        variable_initializer: VariableInitializer,
        streaming_handler: StreamingHandler,
    ) -> None:
        """Initialize builder with shared services.

        Args:
            config: Global configuration
            variable_initializer: Service for initializing variables
            streaming_handler: Handler for streaming updates
        """
        self.config = config
        self.variable_initializer = variable_initializer
        self.streaming_handler = streaming_handler

    async def build_execution_context(
        self,
        flow_config: FlowConfig,
        variables: dict[str, Any] | None,
        include_trace: bool,
        session_id: str | None,
        workspace_slug: str | None,
        thread_id: str | None,
        agent_id: str | None,
    ) -> tuple[ExecutionContext, MCPSessionPool]:
        """Build execution context and create MCP session pool.

        Args:
            flow_config: Flow configuration
            variables: Initial variables
            include_trace: Whether to include execution trace
            session_id: Session identifier
            workspace_slug: Workspace identifier
            thread_id: Thread identifier
            agent_id: Agent identifier

        Returns:
            Tuple of (ExecutionContext, MCPSessionPool)

        Resource Ownership:
            The CALLER is responsible for closing the returned MCPSessionPool.
            This must happen in a try/finally block to ensure cleanup.
        """
        flow_id = flow_config.uuid

        # Initialize variables
        flow_variables = await self.variable_initializer.initialize_variables(
            flow_config, variables
        )

        # Prepare trace log if requested
        trace_log: list[TraceEntry] | None = [] if include_trace else None

        # Create MCP session pool for persistent connections
        # Ownership transfer: Caller must close this pool
        mcp_session_pool = MCPSessionPool(flow_id)

        execution_context = ExecutionContext(
            variables=flow_variables,
            flow_id=flow_id,
            step_index=0,
            step_id="",
            config=self.config,
            streaming_handler=self.streaming_handler,
            trace_log=trace_log,
            mcp_session_pool=mcp_session_pool,
            session_id=session_id,
            thread_id=thread_id,
            workspace_slug=workspace_slug,
            agent_id=agent_id,
        )

        return execution_context, mcp_session_pool
