from moose_lib import Key
from pydantic import BaseModel
 
class MyModel(BaseModel):
    id: Key[str]  # Marks 'id' as a primary key of type string

from pydantic import BaseModel
 
class MyDataModel(BaseModel):
    id: str
    name: str
    count: int

class MooseClient:
    query: QueryClient  # For database queries
    workflow: Optional[WorkflowClient]  # For workflow operations

class IngestionFormat(Enum):
    JSON = "JSON"  # Single JSON object
    JSON_ARRAY = "JSON_ARRAY"  # Array of JSON objects

@dataclass
class ConsumptionApiResult:
    status: int  # HTTP status code
    body: Any    # Response body

class OlapConfig(BaseModel):
    order_by_fields: list[str] = []  # Fields to order by
    deduplicate: bool = False        # Enable deduplication
    engine: Optional[ClickHouseEngines] = None  # Table engine

class StreamConfig(BaseModel):
    parallelism: int = 1
    retention_period: int = 60 * 60 * 24 * 7  # 7 days
    destination: Optional[OlapTable[Any]] = None

class IngestConfig(BaseModel):
    format: Optional[IngestionFormat] = None
    destination: Optional[OlapTable[Any]] = None

class IngestPipelineConfig(BaseModel):
    table: bool | OlapConfig = True
    stream: bool | StreamConfig = True
    ingest: bool | IngestConfig = True

# Basic usage
my_table = OlapTable[UserProfile]("user_profiles")
 
# With configuration
my_table = OlapTable[UserProfile]("user_profiles", OlapConfig(
    order_by_fields=["id", "timestamp"],
    deduplicate=True
))

# Basic usage
my_stream = Stream[UserEvent]("user_events")
 
# With configuration
my_stream = Stream[UserEvent]("user_events", StreamConfig(
    parallelism=3,
    retention_period=86400  # 1 day in seconds
))
 
# Adding transformations
def transform_user_event(event: UserEvent) -> ProfileUpdate:
    return ProfileUpdate(user_id=event.user_id, update_type="event")
 
my_stream.add_transform(profile_stream, transform_user_event)

# Basic usage with destination stream
my_ingest_api = IngestApi[UserEvent]("user_events", IngestConfigWithDestination(
    destination=my_user_event_stream
))

# Basic usage
def get_user_profiles(params: UserQuery) -> list[UserProfile]:
    # Query implementation
    return [UserProfile(...), UserProfile(...)]
 
my_api = ConsumptionApi[UserQuery, list[UserProfile]](
    "get_user_profiles",
    get_user_profiles
)

# Basic usage
pipeline = IngestPipeline[UserEvent]("user_pipeline", IngestPipelineConfig(
    ingest=True,
    stream=True,
    table=True
))
 
# With advanced configuration
pipeline = IngestPipeline[UserEvent]("user_pipeline", IngestPipelineConfig(
    ingest=True,
    stream=StreamConfig(parallelism=3),
    table=OlapConfig(
        order_by_fields=["id", "timestamp"],
        deduplicate=True
    )
))

# Basic usage
view = MaterializedView[UserStatistics](MaterializedViewOptions(
    select_statement="SELECT user_id, COUNT(*) as event_count FROM user_events GROUP BY user_id",
    table_name="user_events",
    materialized_view_name="user_statistics",
    order_by_fields=["user_id"]
))

class ClickHouseEngines(Enum):
    MergeTree = "MergeTree"
    ReplacingMergeTree = "ReplacingMergeTree"
    SummingMergeTree = "SummingMergeTree"
    AggregatingMergeTree = "AggregatingMergeTree"
    CollapsingMergeTree = "CollapsingMergeTree"
    VersionedCollapsingMergeTree = "VersionedCollapsingMergeTree"
    GraphiteMergeTree = "GraphiteMergeTree"

from moose_lib import task
 
@task(retries=5)
def process_data(input_data: dict) -> dict:
    # Process data
    return {
        "task": "process_data",
        "data": {"processed": True}
    }
 
# Can also be used without parameters
@task
def simple_task() -> dict:
    return {
        "task": "simple_task",
        "data": {"completed": True}
    }