from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field

from pygeai_orchestration.core.base.pattern import BasePattern, PatternResult
from pygeai_orchestration.core.exceptions import PatternExecutionError


class CompositionMode(str, Enum):
    SEQUENTIAL = "sequential"
    PARALLEL = "parallel"


class CompositionConfig(BaseModel):
    mode: CompositionMode = Field(
        default=CompositionMode.SEQUENTIAL,
        description="Execution mode for composed patterns"
    )
    stop_on_error: bool = Field(
        default=True,
        description="Stop execution if a pattern fails"
    )
    pass_output: bool = Field(
        default=True,
        description="Pass output of previous pattern as input to next"
    )


class PatternPipeline:
    def __init__(self, config: Optional[CompositionConfig] = None):
        self.config = config or CompositionConfig()
        self._patterns: list[BasePattern] = []
        self._results: list[PatternResult] = []

    def add_pattern(self, pattern: BasePattern) -> "PatternPipeline":
        if not isinstance(pattern, BasePattern):
            raise TypeError(f"Expected BasePattern, got {type(pattern)}")
        self._patterns.append(pattern)
        return self

    def add_patterns(self, *patterns: BasePattern) -> "PatternPipeline":
        for pattern in patterns:
            self.add_pattern(pattern)
        return self

    async def execute(self, initial_input: str, **kwargs) -> PatternResult:
        if not self._patterns:
            raise PatternExecutionError("No patterns in pipeline")

        self._results.clear()
        current_input = initial_input

        if self.config.mode == CompositionMode.SEQUENTIAL:
            return await self._execute_sequential(current_input, **kwargs)
        elif self.config.mode == CompositionMode.PARALLEL:
            return await self._execute_parallel(current_input, **kwargs)
        else:
            raise PatternExecutionError(f"Unsupported composition mode: {self.config.mode}")

    async def _execute_sequential(self, initial_input: str, **kwargs) -> PatternResult:
        current_input = initial_input
        last_result = None

        for i, pattern in enumerate(self._patterns):
            try:
                result = await pattern.execute(current_input, **kwargs)
                self._results.append(result)
                last_result = result

                if self.config.pass_output and result.result:
                    current_input = result.result

            except Exception as e:
                if self.config.stop_on_error:
                    raise PatternExecutionError(f"Pattern {i} ({type(pattern).__name__}) failed: {str(e)}")

                error_result = PatternResult(
                    result="",
                    success=False,
                    error=str(e),
                    metadata={"pattern_index": i, "error": str(e)}
                )
                self._results.append(error_result)
                last_result = error_result

        if last_result is None:
            raise PatternExecutionError("No results produced")

        return PatternResult(
            result=last_result.result,
            success=all(r.success for r in self._results),
            metadata={
                "pipeline_results": [
                    {
                        "index": i,
                        "success": r.success,
                        "result_length": len(r.result) if r.result else 0
                    }
                    for i, r in enumerate(self._results)
                ],
                "total_patterns": len(self._patterns),
                "successful_patterns": sum(1 for r in self._results if r.success)
            }
        )

    async def _execute_parallel(self, initial_input: str, **kwargs) -> PatternResult:
        import asyncio

        tasks = [pattern.execute(initial_input, **kwargs) for pattern in self._patterns]

        try:
            results = await asyncio.gather(*tasks, return_exceptions=True)
        except Exception as e:
            raise PatternExecutionError(f"Parallel execution failed: {str(e)}")

        self._results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                error_result = PatternResult(
                    result="",
                    success=False,
                    error=str(result),
                    metadata={"pattern_index": i, "error": str(result)}
                )
                self._results.append(error_result)
            else:
                self._results.append(result)

        successful_results = [r for r in self._results if r.success]

        if not successful_results:
            raise PatternExecutionError("All patterns failed in parallel execution")

        combined_output = "\n\n".join(r.result for r in successful_results if r.result)

        return PatternResult(
            result=combined_output,
            success=any(r.success for r in self._results),
            metadata={
                "pipeline_results": [
                    {
                        "index": i,
                        "success": r.success,
                        "result_length": len(r.result) if r.result else 0
                    }
                    for i, r in enumerate(self._results)
                ],
                "total_patterns": len(self._patterns),
                "successful_patterns": len(successful_results),
                "mode": "parallel"
            }
        )

    def get_results(self) -> list[PatternResult]:
        return self._results.copy()

    def clear(self) -> None:
        self._patterns.clear()
        self._results.clear()

    def __len__(self) -> int:
        return len(self._patterns)


class PatternComposer:
    @staticmethod
    def sequential(*patterns: BasePattern) -> PatternPipeline:
        config = CompositionConfig(mode=CompositionMode.SEQUENTIAL)
        pipeline = PatternPipeline(config)
        return pipeline.add_patterns(*patterns)

    @staticmethod
    def parallel(*patterns: BasePattern) -> PatternPipeline:
        config = CompositionConfig(mode=CompositionMode.PARALLEL)
        pipeline = PatternPipeline(config)
        return pipeline.add_patterns(*patterns)

    @staticmethod
    def custom(
        mode: CompositionMode,
        stop_on_error: bool = True,
        pass_output: bool = True
    ) -> PatternPipeline:
        config = CompositionConfig(
            mode=mode,
            stop_on_error=stop_on_error,
            pass_output=pass_output
        )
        return PatternPipeline(config)
