from ai_graph.pipeline.base import Pipeline
from ai_graph.step.base import AddKeyStep, DelKeyStep
from ai_graph.step.foreach import ForEachStep

# Create a simple pipeline
pipeline = Pipeline("DataProcessor")

# Add processing steps
pipeline.add_step(AddKeyStep("status", "processing"))
pipeline.add_step(AddKeyStep("timestamp", "2024-01-01"))

# Process data
data = {"input": "some data"}
result = pipeline.process(data)

print(result)
# Output: {'input': 'some data', 'status': 'processing', 'timestamp': '2024-01-01'}

from ai_graph.step.base import BasePipelineStep

class CustomStep(BasePipelineStep):
    def _process_step(self, data):
        # Your processing logic here
        data["custom_field"] = "processed"
        return data

# Use in pipeline
pipeline = Pipeline("CustomPipeline")
pipeline.add_step(CustomStep("MyCustomStep"))

from ai_graph.pipeline.base import Pipeline

# Create pipeline
pipeline = Pipeline("MyPipeline")

# Add steps (method chaining supported)
pipeline.add_step(step1).add_step(step2).add_step(step3)

# Process data
result = pipeline.process(input_data)

from ai_graph.step.foreach import ForEachStep

# Process a list of items
foreach_step = ForEachStep(
    items_key="items",
    results_key="processed_items"
)

# Add sub-processing steps
foreach_step.add_sub_step(AddKeyStep("processed", True))
foreach_step.add_sub_step(AddKeyStep("batch_id", "batch_001"))

# Use in pipeline
pipeline = Pipeline("BatchProcessor")
pipeline.add_step(foreach_step)

data = {"items": [{"id": 1}, {"id": 2}, {"id": 3}]}
result = pipeline.process(data)

from ai_graph.pipeline.base import Pipeline
from ai_graph.step.base import AddKeyStep

class ValidateDataStep(BasePipelineStep):
    def _process_step(self, data):
        if "required_field" not in data:
            data["validation_error"] = "Missing required field"
        else:
            data["validation_status"] = "valid"
        return data

# Create validation pipeline
pipeline = Pipeline("DataValidator")
pipeline.add_step(ValidateDataStep("Validator"))
pipeline.add_step(AddKeyStep("validated_at", "2024-01-01"))

# Process data
data = {"required_field": "value"}
result = pipeline.process(data)

from ai_graph.step.foreach import ForEachStep

class ProcessItemStep(BasePipelineStep):
    def _process_step(self, data):
        # Access current item and iteration index
        current_item = data["_current_item"]
        iteration_index = data["_iteration_index"]

        # Process the item
        data["processed_value"] = current_item * 2
        data["position"] = iteration_index

        return data

# Create batch processing pipeline
batch_processor = ForEachStep(
    items_key="numbers",
    results_key="processed_numbers"
)
batch_processor.add_sub_step(ProcessItemStep("ItemProcessor"))

pipeline = Pipeline("BatchProcessor")
pipeline.add_step(batch_processor)

# Process batch
data = {"numbers": [1, 2, 3, 4, 5]}
result = pipeline.process(data)
# result["processed_numbers"] will contain processed items

# Run a fixed number of iterations
iteration_step = ForEachStep(
    iterations=5,
    results_key="iteration_results"
)

class IterationStep(BasePipelineStep):
    def _process_step(self, data):
        iteration = data["_iteration_index"]
        data["step"] = f"Step {iteration + 1}"
        return data

iteration_step.add_sub_step(IterationStep("Iterator"))

pipeline = Pipeline("IterationPipeline")
pipeline.add_step(iteration_step)

result = pipeline.process({"start": True})

class BasePipelineStep:
    def __init__(self, name: str = None)
    def set_next(self, step: 'BasePipelineStep') -> None
    def process(self, data: Dict[str, Any]) -> Dict[str, Any]
    def _process_step(self, data: Dict[str, Any]) -> Dict[str, Any]  # Override this

class Pipeline:
    def __init__(self, name: str, first_step: Optional[PipelineStep] = None)
    def add_step(self, step: PipelineStep) -> "Pipeline"
    def process(self, data: Dict[str, Any]) -> Dict[str, Any]

class ForEachStep(BasePipelineStep):
    def __init__(
        self,
        items_key: Optional[str] = None,
        iterations: Optional[int] = None,
        results_key: str = "foreach_results",
        name: str = None
    )
    def add_sub_step(self, step: BasePipelineStep) -> "ForEachStep"