Temporal Io

# Order Fulfillment Saga Example
# Complete implementation of a production-ready order saga with Temporal

"""
Order Fulfillment Saga

This example demonstrates a complete order fulfillment process using the saga pattern:

1. Validate Order      -> No compensation (read-only)
2. Reserve Inventory   -> Release Inventory
3. Process Payment     -> Refund Payment
4. Create Shipment     -> Cancel Shipment
5. Send Confirmation   -> No compensation (notification)

Key features:
- Automatic compensation on failure
- Idempotent operations
- Structured error handling
- Query and signal support
- Production-ready patterns
"""

import asyncio
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from enum import Enum
from typing import Any
from uuid import UUID

from temporalio import activity, workflow
from temporalio.client import Client
from temporalio.common import RetryPolicy
from temporalio.exceptions import ApplicationError
from temporalio.worker import Worker

# ============================================================================
# Domain Models
# ============================================================================


class OrderStatus(str, Enum):
    PENDING = "pending"
    VALIDATING = "validating"
    RESERVING_INVENTORY = "reserving_inventory"
    PROCESSING_PAYMENT = "processing_payment"
    CREATING_SHIPMENT = "creating_shipment"
    SENDING_CONFIRMATION = "sending_confirmation"
    COMPLETED = "completed"
    COMPENSATING = "compensating"
    FAILED = "failed"


@dataclass
class OrderItem:
    sku: str
    quantity: int
    unit_price: Decimal


@dataclass
class Address:
    street: str
    city: str
    state: str
    postal_code: str
    country: str


@dataclass
class OrderInput:
    """Input for the order fulfillment workflow."""
    order_id: str
    customer_id: str
    items: list[OrderItem]
    shipping_address: Address
    billing_address: Address
    idempotency_key: str  # For payment idempotency


@dataclass
class OrderResult:
    """Result of the order fulfillment workflow."""
    order_id: str
    status: str
    reservation_id: str | None = None
    payment_id: str | None = None
    shipment_id: str | None = None
    tracking_number: str | None = None
    completed_at: datetime | None = None
    error: str | None = None


# Activity inputs/outputs
@dataclass
class ValidationResult:
    valid: bool
    total_amount: Decimal
    errors: list[str] = field(default_factory=list)


@dataclass
class ReservationResult:
    reservation_id: str
    items_reserved: int
    warehouse_id: str


@dataclass
class PaymentResult:
    payment_id: str
    amount: Decimal
    status: str
    transaction_id: str


@dataclass
class ShipmentResult:
    shipment_id: str
    tracking_number: str
    carrier: str
    estimated_delivery: datetime


# ============================================================================
# Activities
# ============================================================================


@activity.defn
async def validate_order(order: OrderInput) -> ValidationResult:
    """
    Validate order details.
    No compensation needed - read-only operation.
    """
    activity.logger.info(f"Validating order {order.order_id}")
    errors = []

    # Validate items
    if not order.items:
        errors.append("Order must have at least one item")

    for item in order.items:
        if item.quantity <= 0:
            errors.append(f"Invalid quantity for {item.sku}")
        if item.unit_price <= 0:
            errors.append(f"Invalid price for {item.sku}")

    # Calculate total
    total = sum(
        item.unit_price * item.quantity
        for item in order.items
    )

    if total <= 0:
        errors.append("Order total must be positive")

    # Validate addresses
    if not order.shipping_address.postal_code:
        errors.append("Shipping postal code is required")

    return ValidationResult(
        valid=len(errors) == 0,
        total_amount=total,
        errors=errors,
    )


@activity.defn
async def reserve_inventory(order: OrderInput) -> ReservationResult:
    """
    Reserve inventory for order items.
    Compensation: release_inventory
    """
    activity.logger.info(f"Reserving inventory for order {order.order_id}")

    # Simulate inventory API call
    # In production: call inventory service with idempotency key
    await asyncio.sleep(0.5)

    # Simulate occasional failures for testing
    # import random
    # if random.random() < 0.1:
    #     raise ApplicationError("Inventory service unavailable", non_retryable=False)

    return ReservationResult(
        reservation_id=f"res-{order.order_id}",
        items_reserved=len(order.items),
        warehouse_id="WH-001",
    )


@activity.defn
async def release_inventory(reservation_id: str) -> None:
    """
    Compensate: Release reserved inventory.
    Must be idempotent - safe to call multiple times.
    """
    activity.logger.info(f"Releasing inventory reservation {reservation_id}")

    # Simulate release API call
    await asyncio.sleep(0.2)

    # Idempotent: releasing non-existent reservation is a no-op
    activity.logger.info(f"Inventory released: {reservation_id}")


@activity.defn
async def process_payment(
    order: OrderInput,
    amount: Decimal,
) -> PaymentResult:
    """
    Process customer payment.
    Compensation: refund_payment

    Uses idempotency_key to prevent duplicate charges on retry.
    """
    activity.logger.info(
        f"Processing payment for order {order.order_id}, amount: {amount}"
    )

    # Simulate payment API call with idempotency key
    await asyncio.sleep(1.0)

    # Example: simulate payment declined
    # if amount > 1000:
    #     raise ApplicationError(
    #         "Payment declined: insufficient funds",
    #         non_retryable=True,
    #         type="PaymentDeclined",
    #     )

    return PaymentResult(
        payment_id=f"pay-{order.order_id}",
        amount=amount,
        status="captured",
        transaction_id=f"txn-{order.idempotency_key}",
    )


@activity.defn
async def refund_payment(payment_id: str, amount: Decimal) -> None:
    """
    Compensate: Refund payment.
    Must be idempotent - safe to call multiple times.
    """
    activity.logger.info(f"Refunding payment {payment_id}, amount: {amount}")

    # Simulate refund API call
    await asyncio.sleep(0.5)

    # Idempotent: refunding already-refunded payment is a no-op
    activity.logger.info(f"Payment refunded: {payment_id}")


@activity.defn
async def create_shipment(
    order: OrderInput,
    reservation_id: str,
) -> ShipmentResult:
    """
    Create shipment for the order.
    Compensation: cancel_shipment
    """
    activity.logger.info(f"Creating shipment for order {order.order_id}")

    # Simulate shipping API call
    await asyncio.sleep(0.5)

    return ShipmentResult(
        shipment_id=f"ship-{order.order_id}",
        tracking_number=f"TRK-{order.order_id}-001",
        carrier="FedEx",
        estimated_delivery=datetime.now(timezone.utc) + timedelta(days=3),
    )


@activity.defn
async def cancel_shipment(shipment_id: str) -> None:
    """
    Compensate: Cancel shipment.
    Must be idempotent - safe to call multiple times.
    """
    activity.logger.info(f"Cancelling shipment {shipment_id}")

    # Simulate cancel API call
    await asyncio.sleep(0.2)

    # Idempotent: cancelling already-cancelled shipment is a no-op
    activity.logger.info(f"Shipment cancelled: {shipment_id}")


@activity.defn
async def send_order_confirmation(
    order: OrderInput,
    payment: PaymentResult,
    shipment: ShipmentResult,
) -> dict:
    """
    Send order confirmation email.
    No compensation needed - notification is best-effort.
    """
    activity.logger.info(f"Sending confirmation for order {order.order_id}")

    # Simulate email send
    await asyncio.sleep(0.2)

    return {
        "notification_id": f"notif-{order.order_id}",
        "sent_to": f"customer-{order.customer_id}@example.com",
        "tracking_number": shipment.tracking_number,
    }


# ============================================================================
# Workflow
# ============================================================================


@workflow.defn
class OrderFulfillmentWorkflow:
    """
    Order fulfillment saga workflow.

    Orchestrates the complete order fulfillment process with
    automatic compensation on failure.
    """

    def __init__(self):
        self._status = OrderStatus.PENDING
        self._order: OrderInput | None = None
        self._validation: ValidationResult | None = None
        self._reservation: ReservationResult | None = None
        self._payment: PaymentResult | None = None
        self._shipment: ShipmentResult | None = None
        self._error: str | None = None
        self._compensations: list[tuple[str, Any]] = []

    @workflow.run
    async def run(self, order: OrderInput) -> OrderResult:
        """Execute the order fulfillment saga."""
        self._order = order
        workflow.logger.info(f"Starting order fulfillment: {order.order_id}")

        try:
            # Step 1: Validate order (no compensation)
            self._status = OrderStatus.VALIDATING
            self._validation = await workflow.execute_activity(
                validate_order,
                order,
                start_to_close_timeout=timedelta(seconds=30),
            )

            if not self._validation.valid:
                raise ApplicationError(
                    f"Order validation failed: {self._validation.errors}",
                    non_retryable=True,
                    type="ValidationFailed",
                )

            # Step 2: Reserve inventory
            self._status = OrderStatus.RESERVING_INVENTORY
            self._reservation = await workflow.execute_activity(
                reserve_inventory,
                order,
                start_to_close_timeout=timedelta(minutes=2),
                retry_policy=RetryPolicy(
                    maximum_attempts=3,
                    initial_interval=timedelta(seconds=1),
                ),
            )
            # Register compensation
            self._compensations.append(
                ("release_inventory", self._reservation.reservation_id)
            )

            # Step 3: Process payment
            self._status = OrderStatus.PROCESSING_PAYMENT
            self._payment = await workflow.execute_activity(
                process_payment,
                args=[order, self._validation.total_amount],
                start_to_close_timeout=timedelta(minutes=5),
                retry_policy=RetryPolicy(
                    maximum_attempts=3,
                    initial_interval=timedelta(seconds=2),
                    non_retryable_error_types=["PaymentDeclined"],
                ),
            )
            # Register compensation
            self._compensations.append(
                ("refund_payment", (self._payment.payment_id, self._payment.amount))
            )

            # Step 4: Create shipment
            self._status = OrderStatus.CREATING_SHIPMENT
            self._shipment = await workflow.execute_activity(
                create_shipment,
                args=[order, self._reservation.reservation_id],
                start_to_close_timeout=timedelta(minutes=3),
                retry_policy=RetryPolicy(maximum_attempts=3),
            )
            # Register compensation
            self._compensations.append(
                ("cancel_shipment", self._shipment.shipment_id)
            )

            # Step 5: Send confirmation (no compensation)
            self._status = OrderStatus.SENDING_CONFIRMATION
            await workflow.execute_activity(
                send_order_confirmation,
                args=[order, self._payment, self._shipment],
                start_to_close_timeout=timedelta(minutes=1),
                retry_policy=RetryPolicy(maximum_attempts=2),
            )

            # Success!
            self._status = OrderStatus.COMPLETED
            workflow.logger.info(f"Order completed: {order.order_id}")

            return OrderResult(
                order_id=order.order_id,
                status="completed",
                reservation_id=self._reservation.reservation_id,
                payment_id=self._payment.payment_id,
                shipment_id=self._shipment.shipment_id,
                tracking_number=self._shipment.tracking_number,
                completed_at=workflow.now(),
            )

        except Exception as e:
            self._error = str(e)
            self._status = OrderStatus.COMPENSATING

            workflow.logger.warning(
                f"Order {order.order_id} failed, running "
                f"{len(self._compensations)} compensations: {e}"
            )

            # Run compensations in reverse order
            await self._run_compensations()

            self._status = OrderStatus.FAILED

            return OrderResult(
                order_id=order.order_id,
                status="failed",
                reservation_id=self._reservation.reservation_id if self._reservation else None,
                payment_id=self._payment.payment_id if self._payment else None,
                shipment_id=self._shipment.shipment_id if self._shipment else None,
                error=str(e),
            )

    async def _run_compensations(self):
        """Run all registered compensations in reverse order."""
        for compensation_name, compensation_args in reversed(self._compensations):
            try:
                workflow.logger.info(f"Running compensation: {compensation_name}")

                if compensation_name == "release_inventory":
                    await workflow.execute_activity(
                        release_inventory,
                        compensation_args,
                        start_to_close_timeout=timedelta(minutes=2),
                        retry_policy=RetryPolicy(maximum_attempts=5),
                    )
                elif compensation_name == "refund_payment":
                    payment_id, amount = compensation_args
                    await workflow.execute_activity(
                        refund_payment,
                        args=[payment_id, amount],
                        start_to_close_timeout=timedelta(minutes=5),
                        retry_policy=RetryPolicy(maximum_attempts=5),
                    )
                elif compensation_name == "cancel_shipment":
                    await workflow.execute_activity(
                        cancel_shipment,
                        compensation_args,
                        start_to_close_timeout=timedelta(minutes=2),
                        retry_policy=RetryPolicy(maximum_attempts=5),
                    )

            except Exception as comp_error:
                # Log but continue with other compensations
                workflow.logger.error(
                    f"Compensation {compensation_name} failed: {comp_error}. "
                    "Manual intervention required."
                )

    # =========================================================================
    # Queries - Read workflow state
    # =========================================================================

    @workflow.query
    def get_status(self) -> str:
        """Get current order status."""
        return self._status.value

    @workflow.query
    def get_order_details(self) -> dict:
        """Get full order details."""
        return {
            "order_id": self._order.order_id if self._order else None,
            "status": self._status.value,
            "total_amount": str(self._validation.total_amount) if self._validation else None,
            "reservation_id": self._reservation.reservation_id if self._reservation else None,
            "payment_id": self._payment.payment_id if self._payment else None,
            "shipment_id": self._shipment.shipment_id if self._shipment else None,
            "tracking_number": self._shipment.tracking_number if self._shipment else None,
            "error": self._error,
        }

    # =========================================================================
    # Signals - External input to workflow
    # =========================================================================

    @workflow.signal
    async def expedite_shipping(self):
        """
        Signal to expedite shipping.
        Only valid before shipment is created.
        """
        if self._shipment:
            workflow.logger.warning("Cannot expedite - already shipped")
            return

        workflow.logger.info("Expedite shipping requested")
        # In production: set a flag that create_shipment reads


# ============================================================================
# Client Usage Example
# ============================================================================


async def example_usage():
    """Example: Starting and monitoring order workflow."""

    # Connect to Temporal
    client = await Client.connect("localhost:7233")

    # Create order input
    order = OrderInput(
        order_id="ORD-12345",
        customer_id="CUST-001",
        items=[
            OrderItem(sku="SKU-001", quantity=2, unit_price=Decimal("29.99")),
            OrderItem(sku="SKU-002", quantity=1, unit_price=Decimal("49.99")),
        ],
        shipping_address=Address(
            street="123 Main St",
            city="San Francisco",
            state="CA",
            postal_code="94105",
            country="US",
        ),
        billing_address=Address(
            street="123 Main St",
            city="San Francisco",
            state="CA",
            postal_code="94105",
            country="US",
        ),
        idempotency_key="idem-ORD-12345-v1",
    )

    # Start workflow
    handle = await client.start_workflow(
        OrderFulfillmentWorkflow.run,
        order,
        id=f"order-{order.order_id}",
        task_queue="order-processing",
    )

    print(f"Started workflow: {handle.id}")

    # Query status while running
    await asyncio.sleep(1)
    status = await handle.query(OrderFulfillmentWorkflow.get_status)
    print(f"Current status: {status}")

    # Wait for completion
    result = await handle.result()
    print(f"Order completed: {result}")

    return result


# ============================================================================
# Worker Setup
# ============================================================================


async def run_worker():
    """Run the order fulfillment worker."""
    client = await Client.connect("localhost:7233")

    worker = Worker(
        client,
        task_queue="order-processing",
        workflows=[OrderFulfillmentWorkflow],
        activities=[
            validate_order,
            reserve_inventory,
            release_inventory,
            process_payment,
            refund_payment,
            create_shipment,
            cancel_shipment,
            send_order_confirmation,
        ],
    )

    print("Starting order fulfillment worker...")
    await worker.run()


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "client":
        asyncio.run(example_usage())
    else:
        asyncio.run(run_worker())

# Subscription Billing Workflow Example
# Recurring billing with Temporal - production patterns

"""
Subscription Billing Workflow

This example demonstrates a production-ready recurring billing system:

1. Subscription lifecycle management (create, pause, resume, cancel)
2. Recurring billing with retry and grace period
3. Dunning process for failed payments
4. Proration for plan changes
5. Usage-based billing aggregation

Key features:
- Long-running workflow (months/years)
- Continue-as-new for history management
- Signals for lifecycle events
- Queries for subscription status
- Durable timers for billing cycles
"""

import asyncio
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from enum import Enum
from typing import Any

from temporalio import activity, workflow
from temporalio.client import Client
from temporalio.common import RetryPolicy
from temporalio.exceptions import ApplicationError
from temporalio.worker import Worker

# ============================================================================
# Domain Models
# ============================================================================


class SubscriptionStatus(str, Enum):
    ACTIVE = "active"
    PAST_DUE = "past_due"
    PAUSED = "paused"
    CANCELLED = "cancelled"
    EXPIRED = "expired"


class BillingInterval(str, Enum):
    MONTHLY = "monthly"
    QUARTERLY = "quarterly"
    YEARLY = "yearly"


@dataclass
class Plan:
    plan_id: str
    name: str
    amount: Decimal
    interval: BillingInterval
    trial_days: int = 0


@dataclass
class SubscriptionInput:
    """Input for creating a subscription."""
    subscription_id: str
    customer_id: str
    plan: Plan
    payment_method_id: str
    start_date: datetime | None = None  # None = start immediately


@dataclass
class SubscriptionState:
    """Current state of a subscription."""
    subscription_id: str
    customer_id: str
    plan: Plan
    status: SubscriptionStatus
    current_period_start: datetime
    current_period_end: datetime
    payment_method_id: str
    billing_cycle_count: int = 0
    failed_payment_attempts: int = 0
    last_payment_date: datetime | None = None
    cancelled_at: datetime | None = None
    pause_start: datetime | None = None
    metadata: dict = field(default_factory=dict)


@dataclass
class InvoiceResult:
    invoice_id: str
    amount: Decimal
    status: str  # "paid", "failed", "pending"
    payment_id: str | None = None
    error: str | None = None


@dataclass
class UsageRecord:
    """Usage-based billing record."""
    subscription_id: str
    metric: str
    quantity: int
    unit_price: Decimal
    recorded_at: datetime


# ============================================================================
# Activities
# ============================================================================


@activity.defn
async def create_invoice(
    subscription_id: str,
    customer_id: str,
    amount: Decimal,
    period_start: datetime,
    period_end: datetime,
    line_items: list[dict] | None = None,
) -> str:
    """Create an invoice for the billing period."""
    activity.logger.info(
        f"Creating invoice for subscription {subscription_id}, amount: {amount}"
    )

    # Simulate invoice creation in billing system
    await asyncio.sleep(0.2)

    invoice_id = f"inv-{subscription_id}-{period_start.strftime('%Y%m%d')}"
    activity.logger.info(f"Invoice created: {invoice_id}")

    return invoice_id


@activity.defn
async def charge_payment(
    invoice_id: str,
    payment_method_id: str,
    amount: Decimal,
    idempotency_key: str,
) -> InvoiceResult:
    """Attempt to charge the customer's payment method."""
    activity.logger.info(
        f"Charging payment method {payment_method_id} for invoice {invoice_id}"
    )

    # Simulate payment processing
    await asyncio.sleep(0.5)

    # Simulate occasional payment failures for testing dunning
    # import random
    # if random.random() < 0.2:
    #     return InvoiceResult(
    #         invoice_id=invoice_id,
    #         amount=amount,
    #         status="failed",
    #         error="Card declined",
    #     )

    return InvoiceResult(
        invoice_id=invoice_id,
        amount=amount,
        status="paid",
        payment_id=f"pay-{invoice_id}",
    )


@activity.defn
async def send_payment_failed_notification(
    customer_id: str,
    invoice_id: str,
    attempt_number: int,
    next_retry_date: datetime | None,
) -> None:
    """Notify customer of failed payment."""
    activity.logger.info(
        f"Sending payment failed notification to {customer_id}, "
        f"attempt {attempt_number}"
    )

    # Simulate email/SMS notification
    await asyncio.sleep(0.1)


@activity.defn
async def send_subscription_cancelled_notification(
    customer_id: str,
    subscription_id: str,
    reason: str,
) -> None:
    """Notify customer of subscription cancellation."""
    activity.logger.info(
        f"Sending cancellation notification to {customer_id}: {reason}"
    )
    await asyncio.sleep(0.1)


@activity.defn
async def get_usage_records(
    subscription_id: str,
    period_start: datetime,
    period_end: datetime,
) -> list[UsageRecord]:
    """Fetch usage records for usage-based billing."""
    activity.logger.info(
        f"Fetching usage for {subscription_id} from {period_start} to {period_end}"
    )

    # Simulate fetching from usage tracking system
    await asyncio.sleep(0.2)

    # Example usage records
    return [
        UsageRecord(
            subscription_id=subscription_id,
            metric="api_calls",
            quantity=1500,
            unit_price=Decimal("0.001"),
            recorded_at=period_start + timedelta(days=15),
        ),
        UsageRecord(
            subscription_id=subscription_id,
            metric="storage_gb",
            quantity=25,
            unit_price=Decimal("0.10"),
            recorded_at=period_start + timedelta(days=15),
        ),
    ]


@activity.defn
async def calculate_proration(
    old_plan: Plan,
    new_plan: Plan,
    days_remaining: int,
    total_days: int,
) -> Decimal:
    """Calculate proration amount for plan change."""
    activity.logger.info(
        f"Calculating proration: {old_plan.name} -> {new_plan.name}"
    )

    # Credit for unused time on old plan
    old_daily_rate = old_plan.amount / Decimal(total_days)
    credit = old_daily_rate * days_remaining

    # Charge for remaining time on new plan
    new_daily_rate = new_plan.amount / Decimal(total_days)
    charge = new_daily_rate * days_remaining

    proration = charge - credit

    activity.logger.info(f"Proration amount: {proration}")
    return proration


@activity.defn
async def update_subscription_in_database(state: SubscriptionState) -> None:
    """Persist subscription state to database."""
    activity.logger.info(f"Updating subscription {state.subscription_id} in database")
    await asyncio.sleep(0.1)


# ============================================================================
# Billing Workflow
# ============================================================================


@workflow.defn
class SubscriptionBillingWorkflow:
    """
    Long-running subscription billing workflow.

    Handles:
    - Recurring billing cycles
    - Payment retries (dunning)
    - Subscription lifecycle (pause, resume, cancel)
    - Plan changes with proration
    - Usage-based billing
    """

    # Configuration
    MAX_PAYMENT_ATTEMPTS = 4
    PAYMENT_RETRY_INTERVALS = [
        timedelta(days=3),
        timedelta(days=5),
        timedelta(days=7),
    ]
    GRACE_PERIOD_DAYS = 14
    MAX_BILLING_CYCLES_BEFORE_CONTINUE_AS_NEW = 12  # ~1 year for monthly

    def __init__(self):
        self._state: SubscriptionState | None = None
        self._pending_plan_change: Plan | None = None
        self._cancel_requested = False
        self._pause_requested = False
        self._resume_requested = False

    @workflow.run
    async def run(self, input: SubscriptionInput) -> dict:
        """
        Main subscription billing loop.

        Runs until subscription is cancelled or expires.
        Uses continue-as-new to manage history size.
        """
        # Initialize state
        now = workflow.now()
        period_start = input.start_date or now
        period_end = self._calculate_period_end(period_start, input.plan.interval)

        self._state = SubscriptionState(
            subscription_id=input.subscription_id,
            customer_id=input.customer_id,
            plan=input.plan,
            status=SubscriptionStatus.ACTIVE,
            current_period_start=period_start,
            current_period_end=period_end,
            payment_method_id=input.payment_method_id,
        )

        workflow.logger.info(
            f"Starting subscription {input.subscription_id} for {input.customer_id}"
        )

        # Handle trial period
        if input.plan.trial_days > 0:
            workflow.logger.info(f"Trial period: {input.plan.trial_days} days")
            trial_end = period_start + timedelta(days=input.plan.trial_days)
            await self._wait_until(trial_end)

        # Main billing loop
        while self._state.status == SubscriptionStatus.ACTIVE:
            # Check for continue-as-new
            if self._state.billing_cycle_count >= self.MAX_BILLING_CYCLES_BEFORE_CONTINUE_AS_NEW:
                workflow.logger.info("Continuing as new workflow")
                workflow.continue_as_new(self._create_continuation_input())

            # Wait for billing date
            await self._wait_until(self._state.current_period_end)

            # Check if cancelled or paused during wait
            if self._cancel_requested:
                await self._handle_cancellation("customer_requested")
                break

            if self._pause_requested:
                await self._handle_pause()
                continue

            # Handle pending plan change
            if self._pending_plan_change:
                await self._apply_plan_change()

            # Bill the customer
            success = await self._process_billing_cycle()

            if not success:
                # Dunning failed - cancel subscription
                await self._handle_cancellation("payment_failed")
                break

            # Advance to next period
            self._advance_period()

        # Persist final state
        await workflow.execute_activity(
            update_subscription_in_database,
            self._state,
            start_to_close_timeout=timedelta(seconds=30),
        )

        return {
            "subscription_id": self._state.subscription_id,
            "final_status": self._state.status.value,
            "billing_cycles": self._state.billing_cycle_count,
            "cancelled_at": self._state.cancelled_at.isoformat() if self._state.cancelled_at else None,
        }

    async def _process_billing_cycle(self) -> bool:
        """
        Process a single billing cycle.
        Returns True if payment succeeded, False if all retries exhausted.
        """
        workflow.logger.info(
            f"Processing billing cycle {self._state.billing_cycle_count + 1}"
        )

        # Get usage-based charges
        usage_records = await workflow.execute_activity(
            get_usage_records,
            args=[
                self._state.subscription_id,
                self._state.current_period_start,
                self._state.current_period_end,
            ],
            start_to_close_timeout=timedelta(minutes=2),
        )

        # Calculate total amount
        base_amount = self._state.plan.amount
        usage_amount = sum(
            r.quantity * r.unit_price for r in usage_records
        )
        total_amount = base_amount + usage_amount

        # Create invoice
        invoice_id = await workflow.execute_activity(
            create_invoice,
            args=[
                self._state.subscription_id,
                self._state.customer_id,
                total_amount,
                self._state.current_period_start,
                self._state.current_period_end,
                [{"metric": r.metric, "quantity": r.quantity} for r in usage_records],
            ],
            start_to_close_timeout=timedelta(minutes=1),
        )

        # Attempt payment with retries
        for attempt in range(self.MAX_PAYMENT_ATTEMPTS):
            result = await workflow.execute_activity(
                charge_payment,
                args=[
                    invoice_id,
                    self._state.payment_method_id,
                    total_amount,
                    f"{invoice_id}-attempt-{attempt}",
                ],
                start_to_close_timeout=timedelta(minutes=5),
                retry_policy=RetryPolicy(
                    maximum_attempts=2,  # Network retries
                    non_retryable_error_types=["CardDeclined", "InsufficientFunds"],
                ),
            )

            if result.status == "paid":
                self._state.last_payment_date = workflow.now()
                self._state.failed_payment_attempts = 0
                self._state.status = SubscriptionStatus.ACTIVE
                workflow.logger.info(f"Payment succeeded: {result.payment_id}")
                return True

            # Payment failed
            self._state.failed_payment_attempts = attempt + 1
            self._state.status = SubscriptionStatus.PAST_DUE

            workflow.logger.warning(
                f"Payment failed (attempt {attempt + 1}): {result.error}"
            )

            # Notify customer
            next_retry = None
            if attempt < len(self.PAYMENT_RETRY_INTERVALS):
                next_retry = workflow.now() + self.PAYMENT_RETRY_INTERVALS[attempt]

            await workflow.execute_activity(
                send_payment_failed_notification,
                args=[
                    self._state.customer_id,
                    invoice_id,
                    attempt + 1,
                    next_retry,
                ],
                start_to_close_timeout=timedelta(minutes=1),
            )

            # Wait before retry (unless last attempt)
            if attempt < len(self.PAYMENT_RETRY_INTERVALS):
                await asyncio.sleep(self.PAYMENT_RETRY_INTERVALS[attempt].total_seconds())

        # All attempts failed
        workflow.logger.error(
            f"All payment attempts exhausted for {self._state.subscription_id}"
        )
        return False

    async def _handle_cancellation(self, reason: str):
        """Cancel the subscription."""
        workflow.logger.info(f"Cancelling subscription: {reason}")

        self._state.status = SubscriptionStatus.CANCELLED
        self._state.cancelled_at = workflow.now()

        await workflow.execute_activity(
            send_subscription_cancelled_notification,
            args=[self._state.customer_id, self._state.subscription_id, reason],
            start_to_close_timeout=timedelta(minutes=1),
        )

    async def _handle_pause(self):
        """Pause the subscription."""
        workflow.logger.info("Pausing subscription")

        self._state.status = SubscriptionStatus.PAUSED
        self._state.pause_start = workflow.now()
        self._pause_requested = False

        # Wait for resume signal
        await workflow.wait_condition(
            lambda: self._resume_requested or self._cancel_requested
        )

        if self._resume_requested:
            # Resume from pause
            pause_duration = workflow.now() - self._state.pause_start
            self._state.current_period_end += pause_duration  # Extend period
            self._state.status = SubscriptionStatus.ACTIVE
            self._state.pause_start = None
            self._resume_requested = False
            workflow.logger.info("Subscription resumed")

    async def _apply_plan_change(self):
        """Apply pending plan change with proration."""
        new_plan = self._pending_plan_change
        self._pending_plan_change = None

        now = workflow.now()
        days_remaining = (self._state.current_period_end - now).days
        total_days = (
            self._state.current_period_end - self._state.current_period_start
        ).days

        # Calculate proration
        proration = await workflow.execute_activity(
            calculate_proration,
            args=[self._state.plan, new_plan, days_remaining, total_days],
            start_to_close_timeout=timedelta(seconds=30),
        )

        # Charge or credit proration
        if proration > 0:
            invoice_id = await workflow.execute_activity(
                create_invoice,
                args=[
                    self._state.subscription_id,
                    self._state.customer_id,
                    proration,
                    now,
                    self._state.current_period_end,
                    [{"description": "Plan change proration"}],
                ],
                start_to_close_timeout=timedelta(minutes=1),
            )

            await workflow.execute_activity(
                charge_payment,
                args=[
                    invoice_id,
                    self._state.payment_method_id,
                    proration,
                    f"prorate-{self._state.subscription_id}-{now.isoformat()}",
                ],
                start_to_close_timeout=timedelta(minutes=5),
            )

        # Update plan
        old_plan = self._state.plan
        self._state.plan = new_plan
        workflow.logger.info(f"Plan changed: {old_plan.name} -> {new_plan.name}")

    def _advance_period(self):
        """Advance to the next billing period."""
        self._state.current_period_start = self._state.current_period_end
        self._state.current_period_end = self._calculate_period_end(
            self._state.current_period_start,
            self._state.plan.interval,
        )
        self._state.billing_cycle_count += 1

    def _calculate_period_end(
        self,
        start: datetime,
        interval: BillingInterval,
    ) -> datetime:
        """Calculate billing period end date."""
        if interval == BillingInterval.MONTHLY:
            # Add one month
            if start.month == 12:
                return start.replace(year=start.year + 1, month=1)
            return start.replace(month=start.month + 1)
        elif interval == BillingInterval.QUARTERLY:
            # Add three months
            month = start.month + 3
            year = start.year
            if month > 12:
                month -= 12
                year += 1
            return start.replace(year=year, month=month)
        elif interval == BillingInterval.YEARLY:
            return start.replace(year=start.year + 1)
        else:
            raise ValueError(f"Unknown interval: {interval}")

    async def _wait_until(self, target: datetime):
        """Wait until target time, handling signals."""
        while workflow.now() < target:
            remaining = (target - workflow.now()).total_seconds()

            # Wait in chunks to check for signals
            wait_time = min(remaining, 3600)  # Max 1 hour chunks

            try:
                await workflow.wait_condition(
                    lambda: (
                        self._cancel_requested
                        or self._pause_requested
                        or self._resume_requested
                        or self._pending_plan_change is not None
                    ),
                    timeout=timedelta(seconds=wait_time),
                )
                # Signal received - return to main loop
                return
            except asyncio.TimeoutError:
                # No signal, continue waiting
                continue

    def _create_continuation_input(self) -> SubscriptionInput:
        """Create input for continue-as-new."""
        return SubscriptionInput(
            subscription_id=self._state.subscription_id,
            customer_id=self._state.customer_id,
            plan=self._state.plan,
            payment_method_id=self._state.payment_method_id,
            start_date=self._state.current_period_start,
        )

    # =========================================================================
    # Signals
    # =========================================================================

    @workflow.signal
    async def cancel(self, reason: str = "customer_requested"):
        """Cancel the subscription at end of current period."""
        workflow.logger.info(f"Cancel requested: {reason}")
        self._cancel_requested = True

    @workflow.signal
    async def pause(self):
        """Pause the subscription."""
        workflow.logger.info("Pause requested")
        self._pause_requested = True

    @workflow.signal
    async def resume(self):
        """Resume a paused subscription."""
        workflow.logger.info("Resume requested")
        self._resume_requested = True

    @workflow.signal
    async def change_plan(self, new_plan: Plan):
        """Change subscription plan (takes effect next cycle)."""
        workflow.logger.info(f"Plan change requested: {new_plan.name}")
        self._pending_plan_change = new_plan

    @workflow.signal
    async def update_payment_method(self, payment_method_id: str):
        """Update payment method."""
        workflow.logger.info(f"Payment method updated: {payment_method_id}")
        self._state.payment_method_id = payment_method_id

    # =========================================================================
    # Queries
    # =========================================================================

    @workflow.query
    def get_status(self) -> str:
        """Get subscription status."""
        return self._state.status.value if self._state else "unknown"

    @workflow.query
    def get_state(self) -> dict:
        """Get full subscription state."""
        if not self._state:
            return {}

        return {
            "subscription_id": self._state.subscription_id,
            "customer_id": self._state.customer_id,
            "plan": {
                "id": self._state.plan.plan_id,
                "name": self._state.plan.name,
                "amount": str(self._state.plan.amount),
            },
            "status": self._state.status.value,
            "current_period_start": self._state.current_period_start.isoformat(),
            "current_period_end": self._state.current_period_end.isoformat(),
            "billing_cycle_count": self._state.billing_cycle_count,
            "failed_payment_attempts": self._state.failed_payment_attempts,
            "last_payment_date": (
                self._state.last_payment_date.isoformat()
                if self._state.last_payment_date
                else None
            ),
            "pending_plan_change": (
                self._pending_plan_change.name if self._pending_plan_change else None
            ),
        }

    @workflow.query
    def get_next_billing_date(self) -> str | None:
        """Get next billing date."""
        if not self._state or self._state.status != SubscriptionStatus.ACTIVE:
            return None
        return self._state.current_period_end.isoformat()


# ============================================================================
# Client Usage Example
# ============================================================================


async def example_usage():
    """Example: Create and manage a subscription."""
    client = await Client.connect("localhost:7233")

    # Create subscription
    plan = Plan(
        plan_id="pro-monthly",
        name="Pro Monthly",
        amount=Decimal("29.99"),
        interval=BillingInterval.MONTHLY,
        trial_days=14,
    )

    sub_input = SubscriptionInput(
        subscription_id="sub-12345",
        customer_id="cust-001",
        plan=plan,
        payment_method_id="pm-card-123",
    )

    handle = await client.start_workflow(
        SubscriptionBillingWorkflow.run,
        sub_input,
        id=f"subscription-{sub_input.subscription_id}",
        task_queue="billing",
    )

    print(f"Started subscription workflow: {handle.id}")

    # Query status
    await asyncio.sleep(1)
    status = await handle.query(SubscriptionBillingWorkflow.get_status)
    print(f"Subscription status: {status}")

    # Upgrade plan
    new_plan = Plan(
        plan_id="enterprise-monthly",
        name="Enterprise Monthly",
        amount=Decimal("99.99"),
        interval=BillingInterval.MONTHLY,
    )
    await handle.signal(SubscriptionBillingWorkflow.change_plan, new_plan)
    print("Plan upgrade scheduled")

    # Get full state
    state = await handle.query(SubscriptionBillingWorkflow.get_state)
    print(f"Full state: {state}")


# ============================================================================
# Worker Setup
# ============================================================================


async def run_worker():
    """Run the billing worker."""
    client = await Client.connect("localhost:7233")

    worker = Worker(
        client,
        task_queue="billing",
        workflows=[SubscriptionBillingWorkflow],
        activities=[
            create_invoice,
            charge_payment,
            send_payment_failed_notification,
            send_subscription_cancelled_notification,
            get_usage_records,
            calculate_proration,
            update_subscription_in_database,
        ],
    )

    print("Starting billing worker...")
    await worker.run()


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "client":
        asyncio.run(example_usage())
    else:
        asyncio.run(run_worker())

"""
Saga Workflow Template with Automatic Compensation

Implements the Saga pattern for distributed transactions with:
- Ordered step execution
- Automatic compensation on failure
- Compensation retry logic
- Structured logging
"""
from dataclasses import dataclass, field
from datetime import timedelta
from typing import Any, Callable, TypeVar

from temporalio import activity, workflow
from temporalio.common import RetryPolicy
from temporalio.exceptions import ApplicationError


# ============================================================================
# Data Models
# ============================================================================
@dataclass
class SagaStep:
    """Represents a single saga step with its compensation."""
    name: str
    activity: str
    compensation_activity: str
    args: Any = None
    compensation_args: Any = None
    timeout: timedelta = field(default_factory=lambda: timedelta(minutes=5))


@dataclass
class SagaResult:
    """Result of saga execution."""
    success: bool
    completed_steps: list[str]
    failed_step: str | None = None
    error: str | None = None
    compensation_errors: list[str] = field(default_factory=list)


# ============================================================================
# Activities
# ============================================================================
@activity.defn
async def reserve_inventory(order_id: str, items: list[dict]) -> str:
    """Step 1: Reserve inventory."""
    activity.logger.info(f"Reserving inventory for order {order_id}")
    # Implementation: Call inventory service
    return f"reservation-{order_id}"


@activity.defn
async def release_inventory(reservation_id: str) -> None:
    """Compensation: Release reserved inventory."""
    activity.logger.info(f"Releasing reservation {reservation_id}")
    # Implementation: Call inventory service to release


@activity.defn
async def charge_payment(order_id: str, amount: float) -> str:
    """Step 2: Charge payment."""
    activity.logger.info(f"Charging {amount} for order {order_id}")
    # Implementation: Call payment processor
    return f"payment-{order_id}"


@activity.defn
async def refund_payment(payment_id: str) -> None:
    """Compensation: Refund payment."""
    activity.logger.info(f"Refunding payment {payment_id}")
    # Implementation: Call payment processor to refund


@activity.defn
async def create_shipment(order_id: str, address: dict) -> str:
    """Step 3: Create shipment."""
    activity.logger.info(f"Creating shipment for order {order_id}")
    # Implementation: Call shipping service
    return f"shipment-{order_id}"


@activity.defn
async def cancel_shipment(shipment_id: str) -> None:
    """Compensation: Cancel shipment."""
    activity.logger.info(f"Cancelling shipment {shipment_id}")
    # Implementation: Call shipping service to cancel


# ============================================================================
# Saga Workflow
# ============================================================================
@workflow.defn
class SagaWorkflow:
    """
    Generic Saga workflow with automatic compensation.

    Usage:
        result = await client.execute_workflow(
            SagaWorkflow.run,
            OrderSagaInput(order_id="123", items=[...], amount=99.99),
            id="order-saga-123",
            task_queue="orders",
        )
    """

    def __init__(self):
        self._completed_steps: list[tuple[str, Any]] = []
        self._status = "pending"

    @workflow.run
    async def run(self, input: "OrderSagaInput") -> SagaResult:
        """Execute saga steps with compensation on failure."""
        self._status = "running"

        # Define saga steps in order
        steps = [
            SagaStep(
                name="reserve_inventory",
                activity="reserve_inventory",
                compensation_activity="release_inventory",
                args=(input.order_id, input.items),
                timeout=timedelta(minutes=2),
            ),
            SagaStep(
                name="charge_payment",
                activity="charge_payment",
                compensation_activity="refund_payment",
                args=(input.order_id, input.amount),
                timeout=timedelta(minutes=5),
            ),
            SagaStep(
                name="create_shipment",
                activity="create_shipment",
                compensation_activity="cancel_shipment",
                args=(input.order_id, input.shipping_address),
                timeout=timedelta(minutes=3),
            ),
        ]

        try:
            # Execute steps in order
            for step in steps:
                result = await self._execute_step(step)
                self._completed_steps.append((step.compensation_activity, result))

            self._status = "completed"
            return SagaResult(
                success=True,
                completed_steps=[s.name for s in steps],
            )

        except Exception as e:
            self._status = "compensating"
            workflow.logger.warning(f"Saga failed at step, running compensations: {e}")

            # Run compensations in reverse order
            compensation_errors = await self._run_compensations()

            self._status = "failed"
            return SagaResult(
                success=False,
                completed_steps=[s[0] for s in self._completed_steps],
                failed_step=steps[len(self._completed_steps)].name if self._completed_steps else steps[0].name,
                error=str(e),
                compensation_errors=compensation_errors,
            )

    async def _execute_step(self, step: SagaStep) -> Any:
        """Execute a single saga step."""
        # Map activity names to functions
        activities = {
            "reserve_inventory": reserve_inventory,
            "charge_payment": charge_payment,
            "create_shipment": create_shipment,
        }

        activity_fn = activities[step.activity]
        return await workflow.execute_activity(
            activity_fn,
            args=step.args,
            start_to_close_timeout=step.timeout,
            retry_policy=RetryPolicy(
                maximum_attempts=3,
                initial_interval=timedelta(seconds=1),
                backoff_coefficient=2.0,
            ),
        )

    async def _run_compensations(self) -> list[str]:
        """Run all compensations in reverse order."""
        errors = []

        # Map compensation activity names to functions
        compensations = {
            "release_inventory": release_inventory,
            "refund_payment": refund_payment,
            "cancel_shipment": cancel_shipment,
        }

        for compensation_activity, result in reversed(self._completed_steps):
            try:
                activity_fn = compensations[compensation_activity]
                await workflow.execute_activity(
                    activity_fn,
                    result,
                    start_to_close_timeout=timedelta(minutes=2),
                    retry_policy=RetryPolicy(
                        maximum_attempts=5,  # More retries for compensation
                        initial_interval=timedelta(seconds=2),
                        backoff_coefficient=2.0,
                    ),
                )
            except Exception as e:
                error_msg = f"Compensation {compensation_activity} failed: {e}"
                workflow.logger.error(error_msg)
                errors.append(error_msg)

        return errors

    @workflow.query
    def get_status(self) -> str:
        return self._status

    @workflow.query
    def get_completed_steps(self) -> list[str]:
        return [s[0] for s in self._completed_steps]


# ============================================================================
# Input Model
# ============================================================================
@dataclass
class OrderSagaInput:
    order_id: str
    items: list[dict]
    amount: float
    shipping_address: dict

"""
Scheduled/Cron Workflow Template

Implements recurring workflows with:
- Durable timers (survives worker restarts)
- Continue-as-new for unbounded execution
- Configurable schedule
- Proper error handling
"""
import asyncio
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from enum import Enum
from typing import Any

from temporalio import activity, workflow
from temporalio.client import Client, Schedule, ScheduleActionStartWorkflow, ScheduleSpec, ScheduleIntervalSpec
from temporalio.common import RetryPolicy


# ============================================================================
# Schedule Types
# ============================================================================
class ScheduleType(Enum):
    INTERVAL = "interval"  # Every N minutes/hours
    CRON = "cron"          # Cron expression
    DAILY = "daily"        # Once per day at specific time


@dataclass
class ScheduleConfig:
    schedule_type: ScheduleType
    interval_seconds: int = 3600  # For INTERVAL type
    cron_expression: str = "0 0 * * *"  # For CRON type
    timezone: str = "UTC"


@dataclass
class ScheduledTaskInput:
    task_name: str
    config: dict[str, Any]
    schedule: ScheduleConfig


@dataclass
class ScheduledTaskResult:
    task_name: str
    execution_time: str
    success: bool
    result: Any = None
    error: str | None = None


# ============================================================================
# Activities
# ============================================================================
@activity.defn
async def execute_scheduled_task(task_name: str, config: dict) -> dict:
    """Execute the actual scheduled task."""
    activity.logger.info(f"Executing scheduled task: {task_name}")

    # Heartbeat for long tasks
    activity.heartbeat(f"Starting {task_name}")

    # Replace with your task implementation
    result = {
        "task": task_name,
        "executed_at": datetime.now(timezone.utc).isoformat(),
        "config": config,
    }

    activity.heartbeat(f"Completed {task_name}")
    return result


@activity.defn
async def send_task_notification(task_name: str, success: bool, error: str | None) -> None:
    """Send notification about task completion."""
    activity.logger.info(f"Task {task_name} completed: success={success}")
    # Implementation: Send Slack/email notification


# ============================================================================
# Scheduled Workflow (Internal Loop)
# ============================================================================
@workflow.defn
class ScheduledWorkflow:
    """
    Self-scheduling workflow using durable timers.
    Uses continue-as-new to prevent unbounded history.
    """

    def __init__(self):
        self._execution_count = 0
        self._last_result: ScheduledTaskResult | None = None

    @workflow.run
    async def run(self, input: ScheduledTaskInput) -> None:
        """Run scheduled task in a loop with continue-as-new."""
        max_iterations = 100  # Continue-as-new after 100 iterations

        while self._execution_count < max_iterations:
            # Execute the task
            result = await self._execute_task(input)
            self._last_result = result
            self._execution_count += 1

            # Notify on failure
            if not result.success:
                await workflow.execute_activity(
                    send_task_notification,
                    args=[input.task_name, False, result.error],
                    start_to_close_timeout=timedelta(seconds=30),
                )

            # Sleep until next execution (durable timer)
            await asyncio.sleep(input.schedule.interval_seconds)

        # Continue-as-new to reset history
        workflow.continue_as_new(input)

    async def _execute_task(self, input: ScheduledTaskInput) -> ScheduledTaskResult:
        """Execute single task iteration."""
        try:
            result = await workflow.execute_activity(
                execute_scheduled_task,
                args=[input.task_name, input.config],
                start_to_close_timeout=timedelta(minutes=10),
                heartbeat_timeout=timedelta(minutes=2),
                retry_policy=RetryPolicy(
                    maximum_attempts=3,
                    initial_interval=timedelta(seconds=5),
                    backoff_coefficient=2.0,
                ),
            )
            return ScheduledTaskResult(
                task_name=input.task_name,
                execution_time=workflow.now().isoformat(),
                success=True,
                result=result,
            )
        except Exception as e:
            return ScheduledTaskResult(
                task_name=input.task_name,
                execution_time=workflow.now().isoformat(),
                success=False,
                error=str(e),
            )

    @workflow.query
    def get_execution_count(self) -> int:
        return self._execution_count

    @workflow.query
    def get_last_result(self) -> ScheduledTaskResult | None:
        return self._last_result


# ============================================================================
# Native Schedule (Temporal Schedules API - Recommended)
# ============================================================================
async def create_schedule(
    client: Client,
    schedule_id: str,
    task_name: str,
    config: dict,
    cron: str = "0 * * * *",  # Every hour
) -> None:
    """
    Create a native Temporal schedule (recommended for production).

    Advantages over workflow-based scheduling:
    - Server-managed, no worker needed between executions
    - Built-in pause/resume/backfill
    - Better visibility in UI
    """
    await client.create_schedule(
        schedule_id,
        Schedule(
            action=ScheduleActionStartWorkflow(
                "ScheduledTaskWorkflow",
                args=[task_name, config],
                id=f"scheduled-{task_name}-{{{{.ScheduleTime.Format `20060102-150405`}}}}",
                task_queue="scheduled-tasks",
            ),
            spec=ScheduleSpec(
                cron_expressions=[cron],
            ),
        ),
    )


async def create_interval_schedule(
    client: Client,
    schedule_id: str,
    task_name: str,
    config: dict,
    interval: timedelta = timedelta(hours=1),
) -> None:
    """Create interval-based schedule."""
    await client.create_schedule(
        schedule_id,
        Schedule(
            action=ScheduleActionStartWorkflow(
                "ScheduledTaskWorkflow",
                args=[task_name, config],
                id=f"scheduled-{task_name}-{{{{.ScheduleTime.Format `20060102-150405`}}}}",
                task_queue="scheduled-tasks",
            ),
            spec=ScheduleSpec(
                intervals=[ScheduleIntervalSpec(every=interval)],
            ),
        ),
    )


# ============================================================================
# One-Shot Scheduled Workflow (for native schedules)
# ============================================================================
@workflow.defn
class ScheduledTaskWorkflow:
    """Single execution workflow for use with native schedules."""

    @workflow.run
    async def run(self, task_name: str, config: dict) -> ScheduledTaskResult:
        try:
            result = await workflow.execute_activity(
                execute_scheduled_task,
                args=[task_name, config],
                start_to_close_timeout=timedelta(minutes=10),
                heartbeat_timeout=timedelta(minutes=2),
                retry_policy=RetryPolicy(maximum_attempts=3),
            )
            return ScheduledTaskResult(
                task_name=task_name,
                execution_time=workflow.now().isoformat(),
                success=True,
                result=result,
            )
        except Exception as e:
            # Notify on failure
            await workflow.execute_activity(
                send_task_notification,
                args=[task_name, False, str(e)],
                start_to_close_timeout=timedelta(seconds=30),
            )
            return ScheduledTaskResult(
                task_name=task_name,
                execution_time=workflow.now().isoformat(),
                success=False,
                error=str(e),
            )

"""
Production Temporal Worker Template

Usage:
    python worker.py --task-queue orders --namespace production

Features:
    - Graceful shutdown handling
    - Structured logging
    - Health checks
    - Metrics export (optional Prometheus)
"""
import asyncio
import logging
import signal
import sys
from argparse import ArgumentParser
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import timedelta

from temporalio import activity, workflow
from temporalio.client import Client, TLSConfig
from temporalio.worker import Worker
from temporalio.common import RetryPolicy


# ============================================================================
# Configuration
# ============================================================================
@dataclass
class WorkerConfig:
    temporal_host: str = "localhost:7233"
    namespace: str = "default"
    task_queue: str = "default"
    max_concurrent_activities: int = 100
    max_concurrent_workflow_tasks: int = 100
    # TLS for production
    tls_cert_path: str | None = None
    tls_key_path: str | None = None


# ============================================================================
# Activities
# ============================================================================
@activity.defn
async def example_activity(input: str) -> str:
    """Replace with your activity implementation."""
    activity.logger.info(f"Processing: {input}")

    # Heartbeat for long operations
    for i in range(10):
        activity.heartbeat(f"Step {i}/10")
        await asyncio.sleep(0.1)

    return f"Processed: {input}"


# ============================================================================
# Workflows
# ============================================================================
@workflow.defn
class ExampleWorkflow:
    """Replace with your workflow implementation."""

    @workflow.run
    async def run(self, input: str) -> str:
        result = await workflow.execute_activity(
            example_activity,
            input,
            start_to_close_timeout=timedelta(seconds=30),
            retry_policy=RetryPolicy(
                maximum_attempts=3,
                initial_interval=timedelta(seconds=1),
                backoff_coefficient=2.0,
            ),
        )
        return result


# ============================================================================
# Worker Setup
# ============================================================================
@asynccontextmanager
async def create_worker(config: WorkerConfig):
    """Create worker with proper lifecycle management."""
    # TLS config for production
    tls_config = None
    if config.tls_cert_path and config.tls_key_path:
        with open(config.tls_cert_path, "rb") as f:
            cert = f.read()
        with open(config.tls_key_path, "rb") as f:
            key = f.read()
        tls_config = TLSConfig(client_cert=cert, client_private_key=key)

    # Connect to Temporal
    client = await Client.connect(
        config.temporal_host,
        namespace=config.namespace,
        tls=tls_config,
    )

    # Create worker
    worker = Worker(
        client,
        task_queue=config.task_queue,
        workflows=[ExampleWorkflow],
        activities=[example_activity],
        max_concurrent_activities=config.max_concurrent_activities,
        max_concurrent_workflow_tasks=config.max_concurrent_workflow_tasks,
    )

    try:
        yield worker
    finally:
        await client.close()


async def run_worker(config: WorkerConfig):
    """Run worker with graceful shutdown."""
    shutdown_event = asyncio.Event()

    def signal_handler():
        logging.info("Shutdown signal received")
        shutdown_event.set()

    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, signal_handler)

    async with create_worker(config) as worker:
        logging.info(f"Worker started on task queue: {config.task_queue}")

        # Run until shutdown signal
        worker_task = asyncio.create_task(worker.run())
        shutdown_task = asyncio.create_task(shutdown_event.wait())

        done, pending = await asyncio.wait(
            [worker_task, shutdown_task],
            return_when=asyncio.FIRST_COMPLETED,
        )

        # Cancel pending tasks
        for task in pending:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

        logging.info("Worker shutdown complete")


# ============================================================================
# Entry Point
# ============================================================================
def main():
    parser = ArgumentParser(description="Temporal Worker")
    parser.add_argument("--host", default="localhost:7233")
    parser.add_argument("--namespace", default="default")
    parser.add_argument("--task-queue", default="default")
    parser.add_argument("--max-activities", type=int, default=100)
    parser.add_argument("--tls-cert", default=None)
    parser.add_argument("--tls-key", default=None)
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    config = WorkerConfig(
        temporal_host=args.host,
        namespace=args.namespace,
        task_queue=args.task_queue,
        max_concurrent_activities=args.max_activities,
        tls_cert_path=args.tls_cert,
        tls_key_path=args.tls_key,
    )

    asyncio.run(run_worker(config))


if __name__ == "__main__":
    main()