import unittest

from pygeai_orchestration.core.base.pattern import BasePattern, PatternConfig, PatternResult, PatternType
from pygeai_orchestration.core.composition import (
    CompositionConfig,
    CompositionMode,
    PatternComposer,
    PatternPipeline,
)
from pygeai_orchestration.core.exceptions import PatternExecutionError


class MockPattern(BasePattern):
    def __init__(self, name: str, output: str, success: bool = True):
        config = PatternConfig(name=name, pattern_type=PatternType.REFLECTION)
        super().__init__(config)
        self.output = output
        self.success = success
        self.executed = False

    async def execute(self, task: str, **kwargs) -> PatternResult:
        self.executed = True
        if not self.success:
            raise Exception(f"Pattern {self.config.name} failed")
        return PatternResult(
            success=True,
            result=f"{self.output}: {task}",
            metadata={"pattern_name": self.config.name}
        )

    async def _execute_impl(self, task: str, context=None) -> PatternResult:
        return await self.execute(task)

    async def step(self, task: str, context=None):
        return await self.execute(task)


class TestCompositionConfig(unittest.TestCase):
    def test_default_config(self):
        config = CompositionConfig()
        self.assertEqual(config.mode, CompositionMode.SEQUENTIAL)
        self.assertTrue(config.stop_on_error)
        self.assertTrue(config.pass_output)

    def test_custom_config(self):
        config = CompositionConfig(
            mode=CompositionMode.PARALLEL,
            stop_on_error=False,
            pass_output=False
        )
        self.assertEqual(config.mode, CompositionMode.PARALLEL)
        self.assertFalse(config.stop_on_error)
        self.assertFalse(config.pass_output)


class TestPatternPipeline(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.pattern1 = MockPattern("pattern1", "Result1")
        self.pattern2 = MockPattern("pattern2", "Result2")
        self.pattern3 = MockPattern("pattern3", "Result3")

    async def test_pipeline_creation(self):
        pipeline = PatternPipeline()
        self.assertEqual(len(pipeline), 0)

    async def test_add_pattern(self):
        pipeline = PatternPipeline()
        pipeline.add_pattern(self.pattern1)
        self.assertEqual(len(pipeline), 1)

    async def test_add_patterns(self):
        pipeline = PatternPipeline()
        pipeline.add_patterns(self.pattern1, self.pattern2, self.pattern3)
        self.assertEqual(len(pipeline), 3)

    async def test_add_invalid_pattern(self):
        pipeline = PatternPipeline()
        with self.assertRaises(TypeError):
            pipeline.add_pattern("not a pattern")

    async def test_fluent_interface(self):
        pipeline = PatternPipeline()
        result = pipeline.add_pattern(self.pattern1).add_pattern(self.pattern2)
        self.assertIs(result, pipeline)
        self.assertEqual(len(pipeline), 2)

    async def test_sequential_execution(self):
        pipeline = PatternPipeline()
        pipeline.add_patterns(self.pattern1, self.pattern2, self.pattern3)

        result = await pipeline.execute("input")

        self.assertTrue(result.success)
        self.assertTrue(self.pattern1.executed)
        self.assertTrue(self.pattern2.executed)
        self.assertTrue(self.pattern3.executed)

    async def test_sequential_output_passing(self):
        pipeline = PatternPipeline(CompositionConfig(pass_output=True))
        pipeline.add_patterns(self.pattern1, self.pattern2)

        result = await pipeline.execute("input")

        self.assertIn("Result2", result.result)
        self.assertIn("Result1: input", result.result)

    async def test_sequential_no_output_passing(self):
        pipeline = PatternPipeline(CompositionConfig(pass_output=False))
        pipeline.add_patterns(self.pattern1, self.pattern2)

        result = await pipeline.execute("input")

        self.assertIn("Result2: input", result.result)

    async def test_stop_on_error(self):
        failing_pattern = MockPattern("failing", "fail", success=False)
        pipeline = PatternPipeline(CompositionConfig(stop_on_error=True))
        pipeline.add_patterns(self.pattern1, failing_pattern, self.pattern3)

        with self.assertRaises(PatternExecutionError):
            await pipeline.execute("input")

        self.assertTrue(self.pattern1.executed)
        self.assertFalse(self.pattern3.executed)

    async def test_continue_on_error(self):
        failing_pattern = MockPattern("failing", "fail", success=False)
        pipeline = PatternPipeline(CompositionConfig(stop_on_error=False))
        pipeline.add_patterns(self.pattern1, failing_pattern, self.pattern3)

        await pipeline.execute("input")

        self.assertTrue(self.pattern1.executed)
        self.assertTrue(self.pattern3.executed)

    async def test_empty_pipeline(self):
        pipeline = PatternPipeline()

        with self.assertRaises(PatternExecutionError) as cm:
            await pipeline.execute("input")

        self.assertIn("No patterns", str(cm.exception))

    async def test_get_results(self):
        pipeline = PatternPipeline()
        pipeline.add_patterns(self.pattern1, self.pattern2)

        await pipeline.execute("input")
        results = pipeline.get_results()

        self.assertEqual(len(results), 2)
        self.assertTrue(all(isinstance(r, PatternResult) for r in results))

    async def test_clear_pipeline(self):
        pipeline = PatternPipeline()
        pipeline.add_patterns(self.pattern1, self.pattern2)
        await pipeline.execute("input")

        pipeline.clear()

        self.assertEqual(len(pipeline), 0)
        self.assertEqual(len(pipeline.get_results()), 0)

    async def test_pipeline_metadata(self):
        pipeline = PatternPipeline()
        pipeline.add_patterns(self.pattern1, self.pattern2, self.pattern3)

        result = await pipeline.execute("input")

        self.assertIn("pipeline_results", result.metadata)
        self.assertEqual(result.metadata["total_patterns"], 3)
        self.assertEqual(result.metadata["successful_patterns"], 3)

    async def test_parallel_execution(self):
        pipeline = PatternPipeline(CompositionConfig(mode=CompositionMode.PARALLEL))
        pipeline.add_patterns(self.pattern1, self.pattern2, self.pattern3)

        result = await pipeline.execute("input")

        self.assertTrue(result.success)
        self.assertTrue(self.pattern1.executed)
        self.assertTrue(self.pattern2.executed)
        self.assertTrue(self.pattern3.executed)

    async def test_parallel_combined_output(self):
        pipeline = PatternPipeline(CompositionConfig(mode=CompositionMode.PARALLEL))
        pipeline.add_patterns(self.pattern1, self.pattern2)

        result = await pipeline.execute("input")

        self.assertIn("Result1", result.result)
        self.assertIn("Result2", result.result)

    async def test_parallel_with_failure(self):
        failing_pattern = MockPattern("failing", "fail", success=False)
        pipeline = PatternPipeline(CompositionConfig(mode=CompositionMode.PARALLEL))
        pipeline.add_patterns(self.pattern1, failing_pattern, self.pattern2)

        result = await pipeline.execute("input")

        self.assertTrue(result.success)
        self.assertIn("Result1", result.result)
        self.assertIn("Result2", result.result)

    async def test_parallel_all_fail(self):
        failing1 = MockPattern("fail1", "fail", success=False)
        failing2 = MockPattern("fail2", "fail", success=False)
        pipeline = PatternPipeline(CompositionConfig(mode=CompositionMode.PARALLEL))
        pipeline.add_patterns(failing1, failing2)

        with self.assertRaises(PatternExecutionError) as cm:
            await pipeline.execute("input")

        self.assertIn("All patterns failed", str(cm.exception))


class TestPatternComposer(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.pattern1 = MockPattern("pattern1", "Result1")
        self.pattern2 = MockPattern("pattern2", "Result2")

    async def test_sequential_composer(self):
        pipeline = PatternComposer.sequential(self.pattern1, self.pattern2)

        self.assertIsInstance(pipeline, PatternPipeline)
        self.assertEqual(len(pipeline), 2)
        self.assertEqual(pipeline.config.mode, CompositionMode.SEQUENTIAL)

    async def test_parallel_composer(self):
        pipeline = PatternComposer.parallel(self.pattern1, self.pattern2)

        self.assertIsInstance(pipeline, PatternPipeline)
        self.assertEqual(len(pipeline), 2)
        self.assertEqual(pipeline.config.mode, CompositionMode.PARALLEL)

    async def test_custom_composer(self):
        pipeline = PatternComposer.custom(
            mode=CompositionMode.SEQUENTIAL,
            stop_on_error=False,
            pass_output=False
        )

        self.assertIsInstance(pipeline, PatternPipeline)
        self.assertFalse(pipeline.config.stop_on_error)
        self.assertFalse(pipeline.config.pass_output)

    async def test_sequential_execution_via_composer(self):
        pipeline = PatternComposer.sequential(self.pattern1, self.pattern2)
        result = await pipeline.execute("input")

        self.assertTrue(result.success)
        self.assertTrue(self.pattern1.executed)
        self.assertTrue(self.pattern2.executed)

    async def test_parallel_execution_via_composer(self):
        pipeline = PatternComposer.parallel(self.pattern1, self.pattern2)
        result = await pipeline.execute("input")

        self.assertTrue(result.success)
        self.assertTrue(self.pattern1.executed)
        self.assertTrue(self.pattern2.executed)


if __name__ == "__main__":
    unittest.main()
