from enum import Enum

from pydantic import BaseModel, Field

from prefect import flow


class DataProcessingStage(Enum):
    """Different stages of the data processing pipeline"""

    RAW_INGESTION = "raw_ingestion"
    CLEAN = "clean"
    TRANSFORM = "transform"
    AGGREGATE = "aggregate"
    EXPORT = "export"


class ProcessingParameters(BaseModel):
    """Input parameters for the data processing pipeline.

    ## *possible pipeline stage values:*

    Each stage represents a distinct processing phase:

    - RAW_INGESTION: Initial data load from source systems

    - CLEAN: Apply data cleaning and validation rules

    - TRANSFORM: Apply business transformations

    - AGGREGATE: Create summary statistics and rollups

    - EXPORT: Prepare and format data for downstream systems

    You may do any inline markdown formatting, like [links](https://www.prefect.io).

    <br>

    <hr>

    <br>

    <details>
        <summary>Extra Information</summary>

            🤫 aggregate is broken!

    </details>
    """

    stage: DataProcessingStage = Field(
        description="Select which processing stage to execute"
    )
    batch_size: int = Field(
        default=1000, description="Number of records to process in each batch"
    )


@flow(name="data-processing-pipeline")
def run_processing(parameters: ProcessingParameters) -> None:
    """Simple flow demonstrating parameter documentation."""
    print(f"Running {parameters.stage.value} with batch size {parameters.batch_size}")


if __name__ == "__main__":
    run_processing.serve(name="data-processing-service", tags=["data-pipeline"])