"""
Pipeline execution engine for Aperion Legendary agent system.
"""

import asyncio
import json
import logging
import time
from concurrent.futures import Future, ThreadPoolExecutor
from typing import Any, Optional

# Optional monolith-only dependency: doc_context middleware.
# When running standalone (outside the full Aperion stack), the decorator
# becomes a transparent no-op so the executor works without the monolith.
try:
    from stack.aperion.foundation.doc_context_middleware import with_doc_context
except ImportError:  # standalone / aperion-flow mode
    def with_doc_context(task_type: str = ""):
        """No-op decorator when monolith middleware is unavailable."""
        def decorator(fn):
            return fn
        return decorator

from .pipeline import (
    LLMAssistanceConfig,
    LLMAssistanceMode,
    LLMAssistanceResult,
    Pipeline,
    PipelineContext,
    PipelineStage,
    PipelineStatus,
    StageStatus,
)
from .recovery import get_recovery_engine

logger = logging.getLogger(__name__)


class PipelineExecutionError(Exception):
    """Exception raised during pipeline execution."""

    pass


class LLMAssistanceHelper:
    """Helper class for LLM-powered pipeline stage assistance."""

    def __init__(self):
        """Initialize LLM assistance helper."""
        self._router = None
        self._cache: dict[str, dict[str, Any]] = {}

    def _get_router(self):
        """Lazy load LLM router."""
        if self._router is None:
            try:
                from stack.aperion.foundation.llm.router import get_router

                self._router = get_router()
            except ImportError:
                logger.warning("LLM router not available, assistance disabled")
        return self._router

    def assist_stage(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        config: LLMAssistanceConfig,
    ) -> LLMAssistanceResult:
        """
        Provide LLM assistance for a pipeline stage.

        Args:
            stage: Pipeline stage to assist
            context: Current pipeline context
            config: LLM assistance configuration

        Returns:
            Result of LLM assistance operation
        """
        start_time = time.time()

        if not config.enabled:
            return LLMAssistanceResult(
                success=False,
                mode=config.mode,
                result={},
                provider_used="none",
                elapsed_time=0.0,
                error="LLM assistance disabled",
            )

        # Check cache
        cache_key = self._get_cache_key(stage, context, config)
        if config.cache_results and cache_key in self._cache:
            cached_result = self._cache[cache_key]
            return LLMAssistanceResult(
                success=True,
                mode=config.mode,
                result=cached_result,
                provider_used="cache",
                elapsed_time=time.time() - start_time,
                cached=True,
            )

        router = self._get_router()
        if not router:
            return LLMAssistanceResult(
                success=False,
                mode=config.mode,
                result={},
                provider_used="none",
                elapsed_time=time.time() - start_time,
                error="LLM router not available",
            )

        try:
            # Route based on mode
            from stack.aperion.foundation.llm.router import TaskType

            task_type_map = {
                LLMAssistanceMode.PRE_VALIDATE: TaskType.ERROR_DIAGNOSIS,
                LLMAssistanceMode.ENRICH_CONTEXT: TaskType.DOC_GENERATION,
                LLMAssistanceMode.ANALYZE_OUTPUT: TaskType.DOC_GENERATION,
                LLMAssistanceMode.ERROR_DIAGNOSIS: TaskType.ERROR_DIAGNOSIS,
            }

            task_type = task_type_map.get(config.mode, TaskType.DOC_GENERATION)

            # Get provider
            provider, decision = router.get_provider(task_type)

            # Build prompt based on mode
            prompt = self._build_prompt(stage, context, config)

            # Execute with timeout
            try:
                response_data = provider.chat(prompt)

                # Extract result
                if isinstance(response_data, dict) and "replies" in response_data:
                    result_text = (
                        response_data["replies"][0] if response_data["replies"] else ""
                    )
                else:
                    result_text = str(response_data)

                result = {"text": result_text, "raw": response_data}

                # Cache result
                if config.cache_results:
                    self._cache[cache_key] = result

                # Track usage
                token_estimate = int(len(result_text.split()) * 1.3)
                router.track_usage(
                    decision.provider_name,
                    tokens=token_estimate,
                    elapsed_time=time.time() - start_time,
                    success=True,
                )

                return LLMAssistanceResult(
                    success=True,
                    mode=config.mode,
                    result=result,
                    provider_used=decision.provider_name,
                    elapsed_time=time.time() - start_time,
                )

            except TimeoutError:
                raise

        except Exception as e:
            logger.error(f"LLM assistance failed for stage {stage.id}: {e}")

            if config.fallback_on_error:
                return LLMAssistanceResult(
                    success=False,
                    mode=config.mode,
                    result={},
                    provider_used="none",
                    elapsed_time=time.time() - start_time,
                    error=str(e),
                )
            else:
                raise

    def _build_prompt(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        config: LLMAssistanceConfig,
    ) -> str:
        """Build LLM prompt based on assistance mode."""
        if config.validation_prompt:
            return config.validation_prompt

        if config.mode == LLMAssistanceMode.PRE_VALIDATE:
            return f"""Validate pipeline stage inputs before execution.

Stage: {stage.name}
Context data: {json.dumps(context.data, indent=2)}

Check if:
1. Required data is present
2. Data types are correct
3. Values are within expected ranges
4. No obvious errors or inconsistencies

Respond with:
- "VALID" if everything looks good
- "INVALID: <reason>" if there are issues"""

        elif config.mode == LLMAssistanceMode.ENRICH_CONTEXT:
            return f"""Add insights to pipeline stage context.

Stage: {stage.name}
Context data: {json.dumps(context.data, indent=2)}

Provide:
1. Summary of current state
2. Potential issues or concerns
3. Suggestions for next steps"""

        elif config.mode == LLMAssistanceMode.ANALYZE_OUTPUT:
            return f"""Analyze pipeline stage output.

Stage: {stage.name}
Output data: {json.dumps(context.data, indent=2)}

Provide:
1. Summary of what happened
2. Quality assessment
3. Recommendations"""

        elif config.mode == LLMAssistanceMode.ERROR_DIAGNOSIS:
            error_info = stage.error or "Unknown error"
            return f"""Diagnose pipeline stage failure.

Stage: {stage.name}
Error: {error_info}
Context: {json.dumps(context.data, indent=2)}

Provide:
1. Root cause analysis
2. Possible fixes
3. Prevention strategies"""

        return ""

    def _get_cache_key(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        config: LLMAssistanceConfig,
    ) -> str:
        """Generate cache key for LLM assistance result."""
        context_hash = hash(json.dumps(context.data, sort_keys=True))
        return f"{config.mode.value}_{stage.id}_{context_hash}"


class PipelineExecutor:
    """
    Core pipeline execution engine.

    Handles pipeline scheduling, stage execution, dependency management,
    and error handling for agent workflows.
    """

    def __init__(self, max_workers: int = 4, default_timeout: int = 300):
        """
        Initialize the pipeline executor.

        Args:
            max_workers: Maximum number of concurrent stage executions
            default_timeout: Default timeout for stage execution in seconds
        """
        self.max_workers = max_workers
        self.default_timeout = default_timeout
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.running_pipelines: dict[str, Pipeline] = {}
        self.stage_futures: dict[str, Future] = {}
        self.llm_helper = LLMAssistanceHelper()
        self.recovery_engine = get_recovery_engine()

    @with_doc_context(task_type="pipeline")
    def execute_pipeline(self, pipeline: Pipeline, doc_context: dict[str, Any] | None = None) -> Pipeline:
        """
        Execute a pipeline synchronously with automatic documentation context injection.

        Args:
            pipeline: Pipeline to execute
            doc_context: Documentation context (automatically injected by middleware)

        Returns:
            The executed pipeline with updated status

        Raises:
            PipelineExecutionError: If pipeline execution fails
        """
        logger.info(f"Starting execution of pipeline: {pipeline.name} ({pipeline.id})")
        
        # Log documentation context availability
        if doc_context:
            logger.info(
                f"Pipeline {pipeline.id} has {doc_context.get('count', 0)} relevant docs "
                f"(retrieved in {doc_context.get('retrieval_time_ms', 0):.2f}ms)"
            )

        try:
            # Reset pipeline state for fresh execution
            pipeline.reset_for_execution()

            pipeline.status = PipelineStatus.RUNNING
            pipeline.started_at = time.time()
            self.running_pipelines[pipeline.id] = pipeline

            # Execute stages in dependency order
            while not pipeline.is_complete():
                ready_stages = pipeline.get_ready_stages()

                if not ready_stages:
                    # Check for circular dependencies or other issues
                    pending_stages = [
                        stage
                        for stage in pipeline.stages
                        if stage.status == StageStatus.PENDING
                    ]
                    if pending_stages:
                        error_msg = (
                            f"Pipeline stalled - no ready stages but "
                            f"{len(pending_stages)} pending"
                        )
                        logger.error(error_msg)
                        pipeline.status = PipelineStatus.FAILED
                        pipeline.error = error_msg
                        break
                    else:
                        break

                # Execute ready stages
                futures = {}
                for stage in ready_stages:
                    future = self.executor.submit(
                        self._execute_stage, stage, pipeline.context
                    )
                    futures[stage.id] = future
                    self.stage_futures[stage.id] = future

                # Wait for stage completion
                for stage_id, future in futures.items():
                    try:
                        updated_context = future.result(timeout=self.default_timeout)
                        pipeline.context = updated_context
                        logger.info(f"Stage {stage_id} completed successfully")
                    except Exception as e:
                        stage = pipeline.get_stage(stage_id)
                        if stage:
                            stage.status = StageStatus.FAILED
                            stage.error = str(e)
                        logger.error(f"Stage {stage_id} failed: {e}")

                # Clean up completed futures
                for stage_id in futures:
                    self.stage_futures.pop(stage_id, None)

            # Determine final pipeline status
            if pipeline.has_failed_stages():
                pipeline.status = PipelineStatus.FAILED
                pipeline.error = "One or more stages failed"
            else:
                pipeline.status = PipelineStatus.COMPLETED

            pipeline.completed_at = time.time()

            logger.info(
                f"Pipeline {pipeline.name} completed with status: "
                f"{pipeline.status.value} (duration: {pipeline.duration():.2f}s)"
            )

            return pipeline

        except Exception as e:
            logger.error(f"Pipeline execution failed: {e}")
            pipeline.status = PipelineStatus.FAILED
            pipeline.error = str(e)
            pipeline.completed_at = time.time()
            raise PipelineExecutionError(f"Pipeline execution failed: {e}") from e

        finally:
            self.running_pipelines.pop(pipeline.id, None)

    @with_doc_context(task_type="pipeline")
    async def execute_pipeline_async(self, pipeline: Pipeline, doc_context: dict[str, Any] | None = None) -> Pipeline:
        """
        Execute a pipeline asynchronously with automatic documentation context injection.

        Args:
            pipeline: Pipeline to execute

        Returns:
            The executed pipeline with updated status
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.execute_pipeline, pipeline)

    @with_doc_context(task_type="pipeline")
    def _execute_stage(
        self, stage: PipelineStage, context: PipelineContext, doc_context: dict[str, Any] | None = None
    ) -> PipelineContext:
        """
        Execute a single stage with automatic documentation context injection.

        Args:
            stage: Stage to execute
            context: Pipeline context
            doc_context: Documentation context (automatically injected by middleware)

        Returns:
            Updated pipeline context
        """
        logger.info(f"Executing stage: {stage.name} ({stage.id})")
        
        # Log documentation context availability
        if doc_context:
            logger.debug(
                f"Stage {stage.id} has {doc_context.get('count', 0)} relevant docs "
                f"(retrieved in {doc_context.get('retrieval_time_ms', 0):.2f}ms)"
            )

        try:
            # PRE_VALIDATE: Check inputs before execution
            if (
                stage.llm_assistance
                and stage.llm_assistance.mode == LLMAssistanceMode.PRE_VALIDATE
            ):
                result = self.llm_helper.assist_stage(
                    stage, context, stage.llm_assistance
                )
                if result.success:
                    context.add_llm_insight(
                        stage.id,
                        {
                            "phase": "pre_validate",
                            "result": result.result,
                            "provider": result.provider_used,
                        },
                    )
                    # Check if validation failed
                    if "INVALID" in result.result.get("text", ""):
                        logger.warning(
                            f"Stage {stage.id} pre-validation failed: {result.result}"
                        )
                else:
                    logger.warning(
                        f"Stage {stage.id} pre-validation LLM failed: {result.error}"
                    )

            # ENRICH_CONTEXT: Add insights before execution
            if (
                stage.llm_assistance
                and stage.llm_assistance.mode == LLMAssistanceMode.ENRICH_CONTEXT
            ):
                result = self.llm_helper.assist_stage(
                    stage, context, stage.llm_assistance
                )
                if result.success:
                    context.add_llm_insight(
                        stage.id,
                        {
                            "phase": "enrich_context",
                            "result": result.result,
                            "provider": result.provider_used,
                        },
                    )

            # Execute stage with timeout handling
            updated_context = stage.execute(context)

            # ANALYZE_OUTPUT: Review results after execution
            if (
                stage.llm_assistance
                and stage.llm_assistance.mode == LLMAssistanceMode.ANALYZE_OUTPUT
            ):
                result = self.llm_helper.assist_stage(
                    stage, updated_context, stage.llm_assistance
                )
                if result.success:
                    updated_context.add_llm_insight(
                        stage.id,
                        {
                            "phase": "analyze_output",
                            "result": result.result,
                            "provider": result.provider_used,
                        },
                    )

            return updated_context

        except Exception as e:
            logger.error(f"Stage {stage.id} execution failed: {e}")

            # Attempt recovery if configured
            if stage.recovery_config:
                logger.info(f"Attempting recovery for stage {stage.id}")
                recovery_result = self.recovery_engine.recover_stage(
                    stage, context, e, stage.recovery_config
                )

                # Store recovery metadata
                context.metadata[f"recovery_{stage.id}"] = {
                    "attempts": len(recovery_result.attempts),
                    "success": recovery_result.success,
                    "total_time": recovery_result.total_recovery_time,
                    "strategy": recovery_result.strategy_used.value
                    if recovery_result.strategy_used
                    else None,
                }

                if recovery_result.success and recovery_result.context:
                    logger.info(
                        f"Stage {stage.id} recovered after {len(recovery_result.attempts)} attempts"
                    )
                    # Update circuit breaker on success
                    if stage.recovery_config:
                        self.recovery_engine.record_success(stage.id)
                    return recovery_result.context

                else:
                    logger.error(
                        f"Stage {stage.id} recovery failed after {len(recovery_result.attempts)} attempts"
                    )
                    # Update circuit breaker on failure
                    if stage.recovery_config:
                        self.recovery_engine.record_failure(
                            stage.id, stage.recovery_config
                        )

            # ERROR_DIAGNOSIS: Diagnose the failure
            if (
                stage.llm_assistance
                and stage.llm_assistance.mode == LLMAssistanceMode.ERROR_DIAGNOSIS
            ):
                # Set error on stage for diagnosis
                stage.error = str(e)
                result = self.llm_helper.assist_stage(
                    stage, context, stage.llm_assistance
                )
                if result.success:
                    context.add_llm_insight(
                        stage.id,
                        {
                            "phase": "error_diagnosis",
                            "result": result.result,
                            "provider": result.provider_used,
                            "original_error": str(e),
                        },
                    )
                    logger.info(
                        f"Stage {stage.id} error diagnosis: {result.result.get('text', '')[:200]}"
                    )

            raise

    def cancel_pipeline(self, pipeline_id: str) -> bool:
        """
        Cancel a running pipeline.

        Args:
            pipeline_id: ID of pipeline to cancel

        Returns:
            True if pipeline was cancelled, False if not found or not running
        """
        pipeline = self.running_pipelines.get(pipeline_id)
        if not pipeline:
            return False

        logger.info(f"Cancelling pipeline: {pipeline.name} ({pipeline_id})")

        # Cancel running stage futures
        for stage_id, future in list(self.stage_futures.items()):
            if stage_id.startswith(pipeline_id):
                future.cancel()
                self.stage_futures.pop(stage_id, None)

        # Update pipeline status
        pipeline.status = PipelineStatus.CANCELLED
        pipeline.completed_at = time.time()

        # Clean up
        self.running_pipelines.pop(pipeline_id, None)

        return True

    def get_pipeline_status(self, pipeline_id: str) -> Optional[dict[str, Any]]:
        """
        Get the status of a pipeline.

        Args:
            pipeline_id: ID of pipeline

        Returns:
            Pipeline status information or None if not found
        """
        pipeline = self.running_pipelines.get(pipeline_id)
        if pipeline:
            return pipeline.to_dict()
        return None

    def list_running_pipelines(self) -> list[dict[str, Any]]:
        """
        List all currently running pipelines.

        Returns:
            List of pipeline status information
        """
        return [pipeline.to_dict() for pipeline in self.running_pipelines.values()]

    def shutdown(self, wait: bool = True) -> None:
        """
        Shutdown the executor.

        Args:
            wait: Whether to wait for running tasks to complete
        """
        logger.info("Shutting down pipeline executor")

        # Cancel all running pipelines
        for pipeline_id in list(self.running_pipelines.keys()):
            self.cancel_pipeline(pipeline_id)

        # Shutdown thread pool
        self.executor.shutdown(wait=wait)

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.shutdown()
