"""Main FlowExecutor class for orchestrating flow execution."""

from typing import Any

from agent_flows.core.flow_engine import (
    ExecutionResult,
    ExecutorRegistry,
    FlowOrchestrator,
)
from agent_flows.core.flow_engine.context_builder import ExecutionContextBuilder
from agent_flows.core.flow_engine.flow_validator import FlowValidator
from agent_flows.core.flow_engine.result_builder import FlowResultBuilder
from agent_flows.core.flow_engine.strategy import ProductionStrategy, TestStrategy
from agent_flows.core.flow_management import FlowLoader, FlowRegistry
from agent_flows.core.runtime import MCPSessionPool, VariableInterpolator
from agent_flows.exceptions import FlowExecutionError
from agent_flows.models.config import AgentFlowsConfig
from agent_flows.models.execution import (
    ExecutionContext,
    FlowResult,
)
from agent_flows.models.flow import FlowConfig
from agent_flows.models.test import TestManifest
from agent_flows.services.realtimex import (
    CredentialResolver,
    RealTimeXClient,
    VariableInitializer,
)
from agent_flows.services.streaming import StreamingHandler
from agent_flows.utils.config import load_config
from agent_flows.utils.logging import (
    get_logger,
    setup_logging,
)
from agent_flows.utils.test_utils import parse_test_manifest

log = get_logger(__name__)


class FlowExecutor:
    """Main class for executing Agent Flows.

    This class serves as the public API and coordination layer for flow execution.
    It delegates specific responsibilities to specialized components:

    - ExecutionContextBuilder: Prepares execution context and resources
    - FlowResultBuilder: Transforms internal results to public FlowResult objects
    - FlowOrchestrator: Handles the core step-by-step execution loop
    """

    def __init__(
        self,
        config: AgentFlowsConfig | None = None,
        session_id: str | None = None,
        workspace_slug: str | None = None,
        thread_id: str | None = None,
        agent_id: str | None = None,
    ) -> None:
        """Initialize FlowExecutor.

        Args:
            config: AgentFlowsConfig instance with all configuration
            session_id: Optional identifier used for streaming sessions

        Note:
            If config is not provided, configuration is loaded from environment variables.
        """
        # Load configuration with priority: config > environment variables
        self.config = config if config is not None else load_config()

        # Setup logging from configuration
        setup_logging(level=self.config.logging.level, json_format=self.config.logging.json_format)

        # Session ID for streaming context
        self.session_id = session_id

        # Workspace Slug
        self.workspace_slug = workspace_slug

        # Session ID for streaming context
        self.thread_id = thread_id

        # Agent ID for context
        self.agent_id = agent_id

        # Initialize RealTimeX client (shared across all managers)
        self.rtx_client = RealTimeXClient(self.config)

        # Initialize core components with injected client
        self.flow_registry = FlowRegistry(
            client=self.rtx_client,
            flow_ttl=self.config.cache_ttl,
        )
        self.flow_loader = FlowLoader(self.flow_registry)
        self.executor_registry = ExecutorRegistry()
        self.variable_interpolator = VariableInterpolator()
        self.credential_resolver = CredentialResolver(self.rtx_client)
        self.variable_initializer = VariableInitializer(self.config, self.rtx_client)

        # Initialize result builder
        self.result_builder = FlowResultBuilder()

        # Initialize streaming handler
        self.streaming_handler = StreamingHandler(session_id=session_id)

        # Initialize context builder
        self.context_builder = ExecutionContextBuilder(
            config=self.config,
            variable_initializer=self.variable_initializer,
            streaming_handler=self.streaming_handler,
        )

        # Initialize flow validator
        self._flow_validator = FlowValidator(
            flow_loader=self.flow_loader,
            executor_registry=self.executor_registry,
        )

        # Execution state
        self._closed = False

    async def _load_and_prepare(
        self,
        flow_source: FlowConfig | str,
        variables: dict[str, Any] | None,
        include_trace: bool,
    ) -> tuple[FlowConfig, ExecutionContext, MCPSessionPool]:
        """Load flow and prepare execution context.

        Returns:
            Tuple of (flow_config, execution_context, mcp_session_pool)

        Resource Ownership:
            Caller MUST close the returned MCPSessionPool.
        """
        flow_config = await self.flow_loader.load(flow_source)

        execution_context, mcp_session_pool = await self.context_builder.build_execution_context(
            flow_config=flow_config,
            variables=variables,
            include_trace=include_trace,
            session_id=self.session_id,
            workspace_slug=self.workspace_slug,
            thread_id=self.thread_id,
            agent_id=self.agent_id,
        )

        return flow_config, execution_context, mcp_session_pool

    async def _build_and_run_pipeline(
        self,
        flow_config: FlowConfig,
        execution_context: ExecutionContext,
        test_manifest: TestManifest | None,
    ) -> ExecutionResult:
        """Build execution pipeline and run flow orchestrator.

        Pipeline Construction:
            - Test mode: Uses TestStrategy with pin/target handling stages
            - Production: Uses ProductionStrategy with standard stages

        Orchestration:
            - FlowOrchestrator manages step-by-step execution loop
            - Uses ExecutionStateTracker for mutable state
            - Returns immutable ExecutionResult
        """
        # Select strategy based on mode (test vs production)
        strategy = TestStrategy(test_manifest) if test_manifest else ProductionStrategy()

        # Build pipeline with strategy-specific stages
        pipeline = strategy.build_pipeline(
            executor_registry=self.executor_registry,
            credential_resolver=self.credential_resolver,
        )

        # Create orchestrator and execute flow
        flow_runner = FlowOrchestrator(
            pipeline=pipeline,
            streaming_handler=self.streaming_handler,
        )

        return await flow_runner.run(
            flow_config=flow_config,
            execution_context=execution_context,
            test_manifest=test_manifest,
        )

    async def execute_flow(
        self,
        flow_source: FlowConfig | str,
        variables: dict[str, Any] | None = None,
        test_manifest: TestManifest | None = None,
        include_trace: bool = False,
    ) -> FlowResult:
        """Execute a flow from various sources.

        Execution Flow:
            1. Load flow config and prepare execution context
            2. Build pipeline (strategy-based) and run orchestrator
            3. Transform ExecutionResult → FlowResult
            4. Handle errors with proper logging and streaming

        Component Interaction:
            - ContextBuilder: Creates ExecutionContext + MCPSessionPool
            - Strategy: Selects pipeline stages (test vs production)
            - FlowOrchestrator: Runs step-by-step execution loop
            - ResultBuilder: Transforms internal → public result models

        Resource Management:
            - MCPSessionPool created in step 1, closed in finally block
            - Ensures no connection leaks even on errors
        """
        if self._closed:
            raise FlowExecutionError("FlowExecutor has been closed")

        try:
            # Step 1: Load and prepare
            # Creates: flow_config, execution_context, mcp_session_pool
            # ContextBuilder handles variable initialization and trace setup
            (
                flow_config,
                execution_context,
                mcp_session_pool,
            ) = await self._load_and_prepare(
                flow_source=flow_source,
                variables=variables,
                include_trace=include_trace,
            )

            try:
                # Step 2: Execute flow
                # Strategy selects pipeline stages based on test_manifest presence
                # Orchestrator runs execution loop with state tracking
                execution_result = await self._build_and_run_pipeline(
                    flow_config=flow_config,
                    execution_context=execution_context,
                    test_manifest=test_manifest,
                )

                # Step 3: Transform result
                # ResultBuilder converts ExecutionResult → FlowResult
                # Pure transformation, no side effects
                return self.result_builder.build_final_result(
                    execution_result, test_manifest, flow_config.uuid
                )
            finally:
                # Step 4: Resource cleanup
                # Always close MCP sessions, even on errors
                await mcp_session_pool.close_all()

        except Exception as e:
            # Error handling: log, emit streaming event, build error result
            # Determine flow_id for error context
            flow_id = (
                flow_config.uuid
                if "flow_config" in locals() and isinstance(flow_config, FlowConfig)
                else str(flow_source)
            )

            # Side effects: logging and streaming (stay in executor)
            log.error("Flow execution failed", flow_id=flow_id, error=str(e), exc_info=True)
            self.streaming_handler.flow_error(str(e))

            # Pure transformation: build error result
            return self.result_builder.build_error_result(e, flow_source, variables)

    async def test(
        self,
        flow_source: FlowConfig | str,
        test_manifest: str | dict[str, Any] | TestManifest | None = None,
        include_trace: bool = False,
    ) -> FlowResult:
        """Execute a flow in test mode asynchronously.

        Args:
            flow_source: FlowConfig object, UUID string, or file path for the flow under test
            test_manifest: Test manifest as file path, dict, TestManifest instance, or None for full flow execution
            include_trace: Whether to include step-by-step execution trace

        Returns:
            FlowResult with results of the test execution

        Raises:
            FileNotFoundError: If test manifest file path doesn't exist
            ValueError: If test manifest data is invalid
            FlowExecutionError: If flow execution fails
        """
        parsed_manifest = parse_test_manifest(test_manifest)

        if parsed_manifest:
            log.info(
                "Starting test execution with manifest",
                test_name=parsed_manifest.name,
                target_step=parsed_manifest.target,
                pinned_steps=len(parsed_manifest.pins),
            )
            initial_variables = parsed_manifest.initial_variables
        else:
            log.info("Starting test execution without manifest (full flow)")
            initial_variables = None

        return await self.execute_flow(
            flow_source=flow_source,
            variables=initial_variables,
            test_manifest=parsed_manifest,
            include_trace=include_trace,
        )

    async def list_flows(self) -> list[FlowConfig]:
        """List all available flows with full configurations."""
        if self._closed:
            raise FlowExecutionError("FlowExecutor has been closed")

        return await self.flow_registry.list_flows()

    async def get_flow_info(self, flow_id: str) -> FlowConfig:
        """Get detailed information about a flow.

        Args:
            flow_id: UUID of the flow

        Returns:
            FlowConfig object with complete flow information
        """
        if self._closed:
            raise FlowExecutionError("FlowExecutor has been closed")

        return await self.flow_registry.get_flow(flow_id)

    async def validate_flow(self, flow_source: FlowConfig | str) -> dict[str, Any]:
        """Validate a flow configuration from any supported source.

        Args:
            flow_source: FlowConfig object, UUID string, or file path to validate

        Returns:
            Dictionary with validation results in the format:
            {
                "valid": bool,
                "step_validation_summary": [
                    {
                        "step_id": str,
                        "step_type": str,
                        "valid": bool,
                        "messages": [str, ...]
                    },
                    ...
                ]
            }

        Raises:
            FlowExecutionError: If the executor is closed.
        """
        if self._closed:
            raise FlowExecutionError("FlowExecutor has been closed")
        return await self._flow_validator.validate(flow_source)

    async def close(self) -> None:
        """Close the executor and cleanup resources."""
        if not self._closed:
            log.info("Closing FlowExecutor")
            # Close components (no-op for most, they don't own the client)
            await self.flow_registry.close()
            await self.credential_resolver.close()
            await self.variable_initializer.close()
            # Close the shared RTX client (owns HTTP sessions)
            await self.rtx_client.close()
            self._closed = True

    async def __aenter__(self) -> "FlowExecutor":
        """Async context manager entry."""
        return self

    async def __aexit__(self, *_: Any) -> None:
        """Async context manager exit."""
        await self.close()
