"""
Cognite Functions handler for historic validation orchestration.

This function operates in two modes:
1. INITIAL MODE (no orchestration_id provided):
   - Creates orchestration state
   - Creates self-managing scheduled workflow (every 2 min)
   - Triggers all partition functions
   - Returns immediately

2. MONITOR MODE (orchestration_id provided):
   - Loads orchestration state
   - Checks partition statuses
   - Retriggers incomplete partitions (max 3 retries)
   - Updates state with progress
   - Deletes workflow and state when all complete

Supports multiple concurrent view validations via unique orchestration_id.
"""

import time
from datetime import datetime, timedelta, timezone
from typing import Any

from cognite.client import CogniteClient
from cognite.client.data_classes.data_modeling import ViewId
from cognite.client.data_classes.filters import Equals, HasData
from common.timestamp_collection_state import (
    load_timestamp_collection_state,
    mark_timestamp_collection_complete,
    save_timestamp_collection_state,
)

# ==============================================================================
# UTILITY FUNCTIONS
# ==============================================================================


def format_timestamp_for_dms(dt: datetime) -> str:
    """
    Format datetime for DMS with max 3 decimal places (milliseconds).

    DMS requires: YYYY-MM-DDTHH:MM:SS[.millis][Z|timezone]
    where millis has 1-3 decimal digits (not 6 for microseconds).
    """
    # Convert to ISO format with timezone
    iso_str = dt.isoformat()

    # Handle timezone
    if iso_str.endswith("+00:00"):
        iso_str = iso_str.replace("+00:00", "Z")

    # Truncate microseconds to milliseconds (6 decimals → 3 decimals)
    if "." in iso_str:
        parts = iso_str.split(".")
        if len(parts) == 2:
            # Extract fractional seconds and timezone
            fractional_and_tz = parts[1]
            if "Z" in fractional_and_tz:
                fractional = fractional_and_tz.replace("Z", "")
                # Truncate to 3 digits
                fractional = fractional[:3]
                iso_str = f"{parts[0]}.{fractional}Z"
            elif "+" in fractional_and_tz or "-" in fractional_and_tz:
                # Has timezone offset
                for tz_char in ["+", "-"]:
                    if tz_char in fractional_and_tz:
                        fractional, tz = fractional_and_tz.split(tz_char)
                        fractional = fractional[:3]
                        iso_str = f"{parts[0]}.{fractional}{tz_char}{tz}"
                        break

    return iso_str


# ==============================================================================
# STATE MANAGEMENT FUNCTIONS
# ==============================================================================


def load_orchestration_state(
    client: CogniteClient, orchestration_id: str, data_quality_space: str = "dataQuality"
) -> dict | None:
    """Load orchestration state from OrchestrationStateView using SDK."""
    try:
        from cognite.client.data_classes.data_modeling.ids import ViewId
        from cognite.client.data_classes.filters import Equals

        # Query via VIEW using SDK
        view_id = ViewId(data_quality_space, "OrchestrationStateView", "1")
        result = client.data_modeling.instances.list(
            instance_type="node", sources=view_id, filter=Equals(["node", "space"], data_quality_space), limit=100
        )

        # Find the matching orchestration
        for node in result:
            if not node.properties:
                continue

            # Access properties using ViewId as key (not string)
            view_props = node.properties.get(view_id)
            if not view_props:
                continue

            node_orch_id = view_props.get("orchestration_id")
            if node_orch_id == orchestration_id:
                # Return all properties as a flat dict
                return {"externalId": node.external_id, **view_props}

        return None
    except Exception as e:
        print(f"Error loading state: {e}")
        import traceback

        traceback.print_exc()
        return None


def save_orchestration_state(client: CogniteClient, state: dict, data_quality_space: str = "dataQuality"):
    """Save orchestration state to container."""
    try:
        from cognite.client.data_classes.data_modeling.ids import ContainerId
        from cognite.client.data_classes.data_modeling.instances import NodeApply, NodeOrEdgeData

        external_id = state.get("externalId", f"orch_state_{state['orchestration_id']}")

        # Create source with proper SDK structure
        source = NodeOrEdgeData(
            source=ContainerId(data_quality_space, "OrchestrationState"),
            properties={k: v for k, v in state.items() if k != "externalId"},
        )

        node = NodeApply(space=data_quality_space, external_id=external_id, sources=[source])

        client.data_modeling.instances.apply(nodes=[node])
        print(f"Saved state for {state['orchestration_id']}")
    except Exception as e:
        print(f"Error saving state: {e}")
        raise


def delete_orchestration_state(client: CogniteClient, orchestration_id: str, data_quality_space: str = "dataQuality"):
    """Delete orchestration state."""
    try:
        external_id = f"orch_state_{orchestration_id}"
        client.data_modeling.instances.delete(nodes=[(data_quality_space, external_id)])
        print(f"Deleted state for {orchestration_id}")
    except Exception as e:
        print(f"Warning: Could not delete state: {e}")


def delete_function_schedule(client: CogniteClient, view_external_id: str):
    """Delete function schedule for this view."""
    try:
        schedule_name = f"Monitor {view_external_id} validation"
        print(f"  Deleting function schedule: {schedule_name}")

        # Find schedule by name
        schedules = client.functions.schedules.list(limit=-1)
        for schedule in schedules:
            if schedule.name == schedule_name:
                client.functions.schedules.delete(id=schedule.id)
                print(f"  ✓ Function schedule deleted (ID: {schedule.id})")
                return

        print("  Note: Schedule not found (may already be deleted)")
    except Exception as e:
        print(f"  Warning: Could not delete function schedule: {e}")


def create_function_schedule(
    client: CogniteClient, view_external_id: str, orchestration_id: str, config: dict, secrets: dict[str, str]
) -> str:
    """
    Create a function schedule to monitor orchestration progress.

    Uses client credentials from secrets to avoid nested session issues.

    Args:
        client: CogniteClient
        view_external_id: View external ID for naming
        orchestration_id: Unique orchestration ID to monitor
        config: Full config dict to pass to monitoring calls
        secrets: Function secrets containing client credentials

    Returns:
        Schedule external_id
    """
    from cognite.client.data_classes.iam import ClientCredentials

    # Validate secrets are provided
    if not secrets:
        raise ValueError(
            "Secrets are required for creating function schedules. "
            "The function must be deployed with client-id and client-secret in secrets configuration."
        )

    if not secrets.get("client-id") or not secrets.get("client-secret"):
        raise ValueError(
            f"Missing required secrets: client-id and/or client-secret. Available secret keys: {list(secrets.keys())}"
        )

    # Get function ID (required for schedules)
    function_ext_id = config.get("function_external_id", "data-quality-validation")
    function = client.functions.retrieve(external_id=function_ext_id)
    if not function:
        raise ValueError(f"Function '{function_ext_id}' not found")

    # Delete existing schedule if present (for idempotency)
    try:
        # List schedules for this function and find matching one
        schedules = client.functions.schedules.list(function_id=function.id, limit=-1)
        for schedule in schedules:
            if schedule.name == f"Monitor {view_external_id} validation":
                client.functions.schedules.delete(id=schedule.id)
                print(f"  Deleted existing schedule: {schedule.name}")
                break
    except Exception as e:
        print(f"  Could not delete existing schedule: {e}")

    # Create schedule with client credentials from secrets (avoids nested sessions)
    created = client.functions.schedules.create(
        name=f"Monitor {view_external_id} validation",
        function_id=function.id,
        cron_expression="*/2 * * * *",  # Every 2 minutes
        client_credentials=ClientCredentials(
            client_id=secrets.get("client-id"), client_secret=secrets.get("client-secret")
        ),
        description=f"Monitors orchestration {orchestration_id}",
        data={
            "orchestration_id": orchestration_id,
            **config,  # Include all original config
        },
    )

    print(f"  Schedule created: {created.name} (ID: {created.id})")
    print(f"  Cron: {created.cron_expression}")

    return created.name


def create_sync_trigger(
    client: CogniteClient,
    config: dict,
    max_partition_date: str,
    orchestration_id: str,
    job_run_id: str,
    secrets: dict[str, str],
) -> str:
    """
    Create continuous sync trigger for ongoing validation.

    Uses lastUpdatedTime filter to catch all updates after historic cutoff.

    Args:
        client: CogniteClient
        config: Dict with view_space, view_external_id, view_version, instance_space, workflow_external_id
        max_partition_date: Maximum date from partition ranges (ISO format)
        orchestration_id: Unique orchestration ID
        job_run_id: Job run ID to be used by the sync validation (typically same as orchestration_id)

    Returns:
        Sync trigger external_id
    """
    from cognite.client.data_classes.data_modeling import ViewId
    from cognite.client.data_classes.data_modeling.query import NodeResultSetExpression, Select, SourceSelector
    from cognite.client.data_classes.filters import And, Equals, Range
    from cognite.client.data_classes.iam import ClientCredentials
    from cognite.client.data_classes.workflows import (
        WorkflowDataModelingTriggerRule,
        WorkflowTriggerDataModelingQuery,
        WorkflowTriggerUpsert,
    )

    # Validate secrets are provided
    if not secrets:
        raise ValueError(
            "Secrets are required for creating triggers. "
            "The function must be deployed with client-id and client-secret in secrets configuration."
        )

    if not secrets.get("client-id") or not secrets.get("client-secret"):
        raise ValueError(
            f"Missing required secrets: client-id and/or client-secret. Available secret keys: {list(secrets.keys())}"
        )

    view_id = ViewId(config["view_space"], config["view_external_id"], config["view_version"])

    # Filter: lastUpdatedTime >= max_partition_date
    # This catches both new instances AND updates to old instances
    dm_filter = And(
        HasData(views=[view_id]),
        Equals(["node", "space"], config["instance_space"]),
        Range(property=["node", "lastUpdatedTime"], gte=max_partition_date),
    )

    sync_batch_size = config.get("sync_batch_size", 100)
    sync_batch_timeout = config.get("sync_batch_timeout", 300)
    use_cursor_mode = config.get("use_sync_cursor_mode", False)

    if use_cursor_mode:
        # Cursor mode: trigger fires as a lightweight signal — fetch only system metadata,
        # no view properties. Matches orchestrator timestamp collection: Select(sources=[]).
        query_select = {"n": Select(sources=[])}
    else:
        # Normal mode: fetch full view properties so instances arrive ready for validation.
        query_select = {"n": Select(sources=[SourceSelector(source=view_id, properties=[])])}

    data_modeling_query = WorkflowTriggerDataModelingQuery(
        with_={"n": NodeResultSetExpression(filter=dm_filter, limit=sync_batch_size)},
        select=query_select,
    )

    trigger_rule = WorkflowDataModelingTriggerRule(
        data_modeling_query=data_modeling_query, batch_size=sync_batch_size, batch_timeout=sync_batch_timeout
    )

    # Update workflow task payload to include job_run_id
    print(f"  Updating workflow task payload with job_run_id: {job_run_id}")
    try:
        from cognite.client.data_classes.workflows import WorkflowVersionUpsert

        # Retrieve existing workflow version
        workflow_external_id = config["workflow_external_id"]
        # Retrieve workflow version "1" (consistent with workflow deployment)
        latest_version = client.workflows.versions.retrieve(workflow_external_id, version="1")

        if latest_version:
            workflow_def = latest_version.workflow_definition

            # Update the first task's data to include job_run_id
            if workflow_def and workflow_def.tasks and len(workflow_def.tasks) > 0:
                task = workflow_def.tasks[0]
                if hasattr(task, "parameters") and hasattr(task.parameters, "data"):
                    # Add job_run_id to the task data
                    if isinstance(task.parameters.data, dict):
                        task.parameters.data["job_run_id"] = job_run_id
                        task.parameters.data["initial_sync_cutoff"] = str(max_partition_date)
                        print("  ✓ Updated workflow task payload with job_run_id and initial_sync_cutoff")
                    else:
                        print("  Warning: Task data is not a dict, cannot add job_run_id")

                # Upsert the updated workflow version
                # Convert WorkflowDefinition (read) to WorkflowDefinitionUpsert (write) format
                new_version = WorkflowVersionUpsert(
                    workflow_external_id=workflow_external_id,
                    version=latest_version.version,
                    workflow_definition=workflow_def.as_write(),
                )
                client.workflows.versions.upsert(new_version)
                print(f"  ✓ Workflow version {latest_version.version} updated")
            else:
                print("  Warning: Workflow has no tasks, cannot update payload")
        else:
            print(f"  Warning: Workflow {workflow_external_id} not found, cannot update payload")
    except Exception as e:
        print(f"  Warning: Could not update workflow payload: {e}")
        print("  Continuing with trigger creation...")

    # Use consistent trigger ID (without orchestration_id) so upsert replaces old trigger
    trigger_external_id = f"sync-{config['view_external_id']}"

    trigger = WorkflowTriggerUpsert(
        external_id=trigger_external_id,
        trigger_rule=trigger_rule,
        workflow_external_id=config["workflow_external_id"],
        workflow_version="1",
        metadata={
            "description": f"Sync validation for {config['view_external_id']} (post-historic)",
            "view_space": config["view_space"],
            "view_external_id": config["view_external_id"],
            "view_version": config["view_version"],
            "datamodel": f"{config['datamodel_space']}/{config['datamodel_version']}/{config['datamodel_external_id']}",
            "instance_space": config["instance_space"],
            "orchestration_id": orchestration_id,
            "historic_cutoff_date": str(max_partition_date),
            "created_by": "historic_orchestrator",
        },
    )

    # Get credentials for trigger (required by triggers.upsert API)
    # These are passed as function secrets during deployment
    trigger_credentials = ClientCredentials(
        client_id=secrets.get("client-id"),
        client_secret=secrets.get("client-secret"),
    )

    print(f"  Upserting sync trigger: {trigger_external_id}")
    print("  This will replace any existing sync trigger for this view")

    # Upsert new trigger with exponential backoff retry (query can timeout due to load)
    max_attempts = 5
    base_delay = 2  # Start with 2 seconds

    for attempt in range(max_attempts):
        try:
            print(f"  Attempting to create sync trigger (attempt {attempt + 1}/{max_attempts})...")
            client.workflows.triggers.upsert(trigger, trigger_credentials)
            print("  ✓ Sync trigger created successfully")
            return trigger_external_id
        except Exception as e:
            error_msg = str(e)
            is_timeout = "timed out" in error_msg.lower() or "timeout" in error_msg.lower()

            if attempt < max_attempts - 1 and is_timeout:
                # Exponential backoff
                delay = base_delay * (2**attempt)
                print(f"  Query timed out, retrying in {delay}s... (attempt {attempt + 1}/{max_attempts})")
                time.sleep(delay)
            else:
                # Last attempt or non-timeout error
                print(f"  Failed to create sync trigger: {error_msg}")
                raise


def delete_workflow(client: CogniteClient, workflow_external_id: str):
    """Delete workflow and its triggers."""
    try:
        # Delete triggers first - list all and filter
        all_triggers = client.workflows.triggers.list(limit=-1)
        workflow_triggers = [t for t in all_triggers if t.workflow_external_id == workflow_external_id]

        if workflow_triggers:
            print(f"  Deleting {len(workflow_triggers)} triggers...")
            for trigger in workflow_triggers:
                client.workflows.triggers.delete(external_id=trigger.external_id)
                print(f"    ✓ Deleted trigger: {trigger.external_id}")

        # Delete workflow
        print(f"  Deleting workflow: {workflow_external_id}")
        client.workflows.delete(external_id=workflow_external_id)
        print("  ✓ Workflow deleted")
    except Exception as e:
        print(f"  Warning: Could not delete workflow: {e}")


def check_partition_statuses(
    client: CogniteClient, state: dict, max_retries: int = 3, function_external_id: str = "data-quality-validation"
) -> tuple[list, list, list, dict, dict, int, int]:
    """
    Check status of all partitions.

    Args:
        client: CogniteClient
        state: Orchestration state dict
        max_retries: Maximum error retries (not stored in container, passed from config)

    Returns: (completed, failed, in_progress, completed_dict, failed_dict, total_validated, total_violations)
    """
    completed_dict = state.get("completed_partitions", {}).copy()
    failed_dict = state.get("failed_partitions", {}).copy()
    completed = []
    failed = []
    in_progress = []
    total_validated = 0
    total_violations = 0

    partition_call_ids = state.get("partition_call_ids", {})

    for partition_id, call_id in partition_call_ids.items():
        # Check if already completed or failed
        if partition_id in completed_dict:
            completed.append(partition_id)
            # Still need to aggregate metrics from completed partitions
            try:
                call = client.functions.calls.retrieve(call_id=int(call_id), function_external_id=function_external_id)
                if call and call.status == "Completed":
                    response = call.get_response() if hasattr(call, "get_response") else call.response
                    if response and response.get("status") == "completed":
                        # Aggregate metrics from already-completed partition
                        total_validated += response.get("cumulative_validated", response.get("instances_validated", 0))
                        total_violations += response.get("total_violations", 0)
            except Exception as e:
                print(f"    Warning: Could not retrieve metrics for completed partition {partition_id}: {e}")
            continue
        if partition_id in failed_dict:
            failed.append(partition_id)
            continue

        try:
            call = client.functions.calls.retrieve(call_id=int(call_id), function_external_id=function_external_id)

            # Handle case where function was redeployed (call_id no longer valid)
            if call is None:
                print(
                    f"    Partition {partition_id}: Call not found"
                    " (function may have been redeployed), treating as in-progress"
                )
                in_progress.append(partition_id)
                continue

            if call.status == "Completed":
                # Check if it completed successfully, partially, or with error
                response = call.get_response() if hasattr(call, "get_response") else call.response
                if response:
                    response_status = response.get("status")
                    if response_status == "completed":
                        completed_dict[partition_id] = True
                        completed.append(partition_id)
                        print(f"    Partition {partition_id}: Completed successfully")

                        # Aggregate metrics
                        total_validated += response.get("cumulative_validated", response.get("instances_validated", 0))
                        total_violations += response.get("total_violations", 0)
                    elif response_status == "partial":
                        # Partial = timeout, needs retry
                        in_progress.append(partition_id)
                        print(f"    Partition {partition_id}: Partial (timeout), needs retry")

                        # Still aggregate what was processed so far
                        total_validated += response.get("cumulative_validated", response.get("instances_validated", 0))
                        total_violations += response.get("total_violations", 0)
                    elif response_status == "error":
                        # Error = needs retry (unless max retries reached)
                        in_progress.append(partition_id)
                        print(f"    Partition {partition_id}: Error, needs retry")
                    else:
                        # Unknown status - treat as error to be safe
                        print(f"    Partition {partition_id}: Unknown status '{response_status}', treating as error")
                        in_progress.append(partition_id)
                else:
                    # No response at all - treat as error
                    print(f"    Partition {partition_id}: No response, treating as error")
                    in_progress.append(partition_id)
            elif call.status in ["Failed", "Timeout"]:
                # Check retry count
                retry_count = state.get("partition_retry_count", {}).get(partition_id, 0)

                print(f"    Partition {partition_id}: {call.status} (retry {retry_count}/{max_retries})")

                if retry_count >= max_retries:
                    failed_dict[partition_id] = True
                    failed.append(partition_id)
                    print("      Max retries exceeded, marking as failed")
                else:
                    in_progress.append(partition_id)
                    print("      Will retry")
            else:
                # Running or Queued
                in_progress.append(partition_id)
                print(f"    Partition {partition_id}: {call.status}")

        except Exception as e:
            print(f"    Partition {partition_id}: Error checking status - {e}")
            in_progress.append(partition_id)

    return completed, failed, in_progress, completed_dict, failed_dict, total_validated, total_violations


def retrigger_partition(
    client: CogniteClient,
    validation_function_id: str,
    partition_id: str,
    partition_data: dict,
    retry_count: int,
) -> str:
    """Retrigger a partition."""
    print(f"  Retriggering partition {partition_id} (attempt {retry_count + 1})")

    call = client.functions.call(external_id=validation_function_id, data=partition_data, wait=False)

    return call.id


# ==============================================================================
# EXISTING FUNCTIONS (discover_data_distribution, etc.)
# ==============================================================================


def discover_data_distribution(
    client: CogniteClient,
    instance_space: str,
    view_id: ViewId,
    partition_field: str,
) -> dict[str, Any]:
    """
    Discover data distribution using optimal DMS queries (NO SAMPLING).

    Uses 3 API calls, fetches only 2 instances:
    - 1 count aggregation query
    - 1 sorted query (ascending) for min value
    - 1 sorted query (descending) for max value

    Returns dict with:
        - total_count: int
        - min_value: any
        - max_value: any
        - field_type: str (datetime, numeric, string)
    """
    print(f"Discovering data distribution for {partition_field}...")

    from cognite.client.data_classes.filters import Equals

    # Step 1: Get total count via aggregation
    print("  1/3 Getting count...")
    try:
        result = client.data_modeling.instances.aggregate(
            view=view_id,
            aggregates=[{"count": {"property": "externalId"}}],
            filter=Equals(["node", "space"], instance_space),
            instance_type="node",
        )
        total_count = result[0].value if result else 0
        print(f"      Found {total_count:,} instances")
    except Exception as e:
        raise ValueError(f"Failed to count instances: {e!s}") from e

    if total_count == 0:
        raise ValueError("No instances found in the specified space/view")

    # Step 2: Get MIN/MAX values - try aggregate first, fall back to sync for datetime fields
    print("  2/3 Getting min/max values...")

    # Try aggregate (works for numeric fields only)
    use_sync = False
    try:
        result = client.data_modeling.instances.aggregate(
            view=view_id,
            aggregates=[{"min": {"property": partition_field}}],
            filter=Equals(["node", "space"], instance_space),
            instance_type="node",
        )
        min_value = result[0].value if result else None

        result = client.data_modeling.instances.aggregate(
            view=view_id,
            aggregates=[{"max": {"property": partition_field}}],
            filter=Equals(["node", "space"], instance_space),
            instance_type="node",
        )
        max_value = result[0].value if result else None

        if min_value is None or max_value is None:
            raise ValueError("No values returned from aggregation")

        print(f"      Min: {min_value}")
        print(f"      Max: {max_value}")

    except Exception:
        # Aggregate failed - likely a datetime field, use sync instead
        print("      Aggregate failed (likely datetime field), using sync...")
        use_sync = True

    # If aggregate failed, use sync to collect all values
    if use_sync:
        from cognite.client.data_classes.data_modeling.query import (
            NodeResultSetExpression,
            Query,
            Select,
            SourceSelector,
        )
        from cognite.client.data_classes.filters import And, HasData

        print("      Collecting all values via sync with two_phase mode...")
        sync_start = time.time()
        all_values = []
        chunk_count = 0
        current_cursor = None

        while True:
            chunk_count += 1

            query = Query(
                with_={
                    "nodes": NodeResultSetExpression(
                        filter=And(Equals(["node", "space"], instance_space), HasData(views=[view_id])),
                        limit=1000,
                        sync_mode="two_phase",
                    )
                },
                select={"nodes": Select(sources=[SourceSelector(source=view_id, properties=[partition_field])])},
                cursors={"nodes": current_cursor} if current_cursor else None,
            )

            result = client.data_modeling.instances.sync(query=query)

            # Extract field values from nodes
            chunk_size = 0
            if hasattr(result, "data") and "nodes" in result.data:
                for node in result.data["nodes"]:
                    # Access property value through the view
                    if hasattr(node, "properties") and node.properties:
                        props = node.properties.get(view_id)
                        if props and partition_field in props:
                            all_values.append(props[partition_field])
                            chunk_size += 1

            # Progress logging every 10 chunks
            if chunk_count % 10 == 0:
                elapsed = time.time() - sync_start
                rate = len(all_values) / elapsed if elapsed > 0 else 0
                print(
                    f"        Progress: {len(all_values):,} values collected"
                    f" ({rate:.0f}/sec, {chunk_count} chunks, {elapsed:.1f}s)"
                )

            # Break if no data received (backfill phase complete)
            if chunk_size == 0:
                print("        Backfill complete: no more data")
                break

            # Get cursor for next iteration
            new_cursor = None
            if hasattr(result, "cursors"):
                if hasattr(result.cursors, "__dict__"):
                    new_cursor = result.cursors.__dict__.get("nodes")
                elif isinstance(result.cursors, dict):
                    new_cursor = result.cursors.get("nodes")

            has_more = new_cursor is not None and new_cursor != current_cursor
            current_cursor = new_cursor

            if not has_more:
                break

        if not all_values:
            raise ValueError(f"No values found for field {partition_field}")

        all_values.sort()
        min_value = all_values[0]
        max_value = all_values[-1]

        print(f"      Collected {len(all_values):,} values from {chunk_count} chunks")
        print(f"      Min: {min_value}")
        print(f"      Max: {max_value}")

    if min_value is None or max_value is None:
        raise ValueError(f"No values found for field {partition_field}")

    # Determine field type from the values
    field_type = "string"
    if isinstance(min_value, int | float):
        field_type = "numeric"
        # If aggregate returned floats but they're whole numbers, convert to int
        if isinstance(min_value, float) and min_value.is_integer():
            min_value = int(min_value)
        if isinstance(max_value, float) and max_value.is_integer():
            max_value = int(max_value)
    elif isinstance(min_value, str):
        try:
            datetime.fromisoformat(min_value.replace("Z", "+00:00"))
            field_type = "datetime"
        except (ValueError, AttributeError):
            field_type = "string"

    print("  ✓ Distribution discovered (3 queries, 2 instances fetched)")
    print(f"    Total: {total_count:,} | Type: {field_type}")
    print(f"    Range: {min_value} to {max_value}")

    return {
        "total_count": total_count,
        "min_value": min_value,
        "max_value": max_value,
        "field_type": field_type,
    }


def calculate_partition_ranges(
    client: CogniteClient,
    distribution: dict[str, Any],
    partition_count: int,
    instance_space: str,
    view_id: ViewId,
    partition_field: str,
) -> list[dict[str, Any]]:
    """
    Calculate partition ranges optimized for even distribution.

    Strategy:
    1. Create initial equal date/numeric ranges
    2. Count instances in each range using aggregation
    3. Adjust boundaries to balance workload across partitions
    """
    from cognite.client.data_classes.filters import And, Equals
    from cognite.client.data_classes.filters import Range as RangeFilter

    total_count = distribution["total_count"]
    min_value = distribution["min_value"]
    max_value = distribution["max_value"]
    field_type = distribution["field_type"]

    target_per_partition = total_count // partition_count

    print(f"Calculating optimized {partition_count} partition ranges...")
    print(f"  Total instances: {total_count:,}")
    print(f"  Target per partition: ~{target_per_partition:,}")

    # Step 1: Create preliminary equal ranges
    print("  Step 1: Creating preliminary equal ranges...")
    preliminary_ranges = []

    if field_type == "datetime":
        min_dt = datetime.fromisoformat(min_value.replace("Z", "+00:00"))
        max_dt = datetime.fromisoformat(max_value.replace("Z", "+00:00"))

        total_seconds = (max_dt - min_dt).total_seconds()
        seconds_per_partition = total_seconds / partition_count

        for i in range(partition_count):
            start_dt = min_dt + timedelta(seconds=(seconds_per_partition * i))
            end_dt = min_dt + timedelta(seconds=(seconds_per_partition * (i + 1)))

            preliminary_ranges.append(
                {
                    "start": format_timestamp_for_dms(start_dt),
                    "end": format_timestamp_for_dms(end_dt),
                }
            )

    elif field_type == "numeric":
        total_range = max_value - min_value
        range_per_partition = total_range / partition_count

        # Check if values are integers (not floats)
        is_integer_field = isinstance(min_value, int) and isinstance(max_value, int)

        for i in range(partition_count):
            start_val = min_value + (range_per_partition * i)
            end_val = min_value + (range_per_partition * (i + 1))

            # Convert to integers if the field is integer type
            if is_integer_field:
                start_val = int(start_val)
                end_val = int(end_val)
                # For last partition, extend end by +1 to include max value
                # (since range filter uses lt=end, not lte=end)
                if i == partition_count - 1:
                    end_val = end_val + 1

            preliminary_ranges.append(
                {
                    "start": start_val,
                    "end": end_val,
                }
            )
    else:
        raise ValueError(f"Cannot partition by field type: {field_type}")

    # Step 2: Count instances in each preliminary range
    print("  Step 2: Counting instances in each range...")
    range_counts = []

    for i, prange in enumerate(preliminary_ranges):
        try:
            result = client.data_modeling.instances.aggregate(
                view=view_id,
                aggregates=[{"count": {"property": "externalId"}}],
                filter=And(
                    Equals(["node", "space"], instance_space),
                    RangeFilter(
                        property=[view_id.space, f"{view_id.external_id}/{view_id.version}", partition_field],
                        gte=prange["start"],
                        lt=prange["end"],
                    ),
                ),
                instance_type="node",
            )
            count = result[0].value if result else 0
            range_counts.append(count)
            print(f"    Range {i}: {count:,} instances")
        except Exception as e:
            print(f"    Range {i}: Error counting - {e}")
            range_counts.append(target_per_partition)  # Fallback estimate

    # Step 3: Adjust boundaries based on cumulative counts
    print("  Step 3: Adjusting boundaries for even distribution...")

    # Calculate cumulative counts
    cumulative = [0]
    for count in range_counts:
        cumulative.append(cumulative[-1] + count)

    print(f"    Initial distribution: {range_counts}")
    print(f"    Cumulative: {cumulative}")

    # Calculate ideal boundary positions (in terms of instance count)
    ideal_positions = [target_per_partition * (i + 1) for i in range(partition_count - 1)]
    print(f"    Ideal cumulative positions: {ideal_positions}")

    # Adjust boundaries based on where ideal positions fall
    if field_type == "datetime":
        # Parse time boundaries
        min_dt = datetime.fromisoformat(min_value.replace("Z", "+00:00"))
        max_dt = datetime.fromisoformat(max_value.replace("Z", "+00:00"))
        total_seconds = (max_dt - min_dt).total_seconds()

        # Preliminary boundaries in datetime
        preliminary_boundaries_dt = []
        for i in range(partition_count + 1):
            boundary_dt = min_dt + timedelta(seconds=(total_seconds * i / partition_count))
            preliminary_boundaries_dt.append(boundary_dt)

        # Calculate new boundaries
        new_boundaries_dt = [min_dt]  # Start boundary stays

        for ideal_pos in ideal_positions:
            # Find which preliminary range this ideal position falls into
            range_idx = 0
            for i in range(len(cumulative) - 1):
                if cumulative[i] <= ideal_pos < cumulative[i + 1]:
                    range_idx = i
                    break

            # How far into this range (as a fraction)?
            range_start_count = cumulative[range_idx]
            range_end_count = cumulative[range_idx + 1]
            range_count = range_end_count - range_start_count

            if range_count > 0:
                fraction_into_range = (ideal_pos - range_start_count) / range_count
            else:
                fraction_into_range = 0.5  # Default to middle if range is empty

            # Calculate new boundary time
            range_start_time = preliminary_boundaries_dt[range_idx]
            range_end_time = preliminary_boundaries_dt[range_idx + 1]
            range_duration = (range_end_time - range_start_time).total_seconds()

            new_boundary_dt = range_start_time + timedelta(seconds=(range_duration * fraction_into_range))
            new_boundaries_dt.append(new_boundary_dt)

        new_boundaries_dt.append(max_dt)  # End boundary stays

        # Build optimized ranges
        optimized_ranges = []
        for i in range(partition_count):
            optimized_ranges.append(
                {
                    "partition_id": i,
                    "start": format_timestamp_for_dms(new_boundaries_dt[i]),
                    "end": format_timestamp_for_dms(new_boundaries_dt[i + 1]),
                    "estimated_count": target_per_partition,
                }
            )

    elif field_type == "numeric":
        # Similar logic for numeric fields
        total_range = max_value - min_value

        # Check if values are integers
        is_integer_field = isinstance(min_value, int) and isinstance(max_value, int)

        # Preliminary boundaries
        preliminary_boundaries = []
        for i in range(partition_count + 1):
            boundary = min_value + (total_range * i / partition_count)
            # Convert to integer if needed
            if is_integer_field:
                boundary = int(boundary)
            preliminary_boundaries.append(boundary)

        # Calculate new boundaries
        new_boundaries = [min_value]  # Start boundary stays

        for ideal_pos in ideal_positions:
            # Find which preliminary range this ideal position falls into
            range_idx = 0
            for i in range(len(cumulative) - 1):
                if cumulative[i] <= ideal_pos < cumulative[i + 1]:
                    range_idx = i
                    break

            # How far into this range (as a fraction)?
            range_start_count = cumulative[range_idx]
            range_end_count = cumulative[range_idx + 1]
            range_count = range_end_count - range_start_count

            if range_count > 0:
                fraction_into_range = (ideal_pos - range_start_count) / range_count
            else:
                fraction_into_range = 0.5

            # Calculate new boundary value
            range_start_val = preliminary_boundaries[range_idx]
            range_end_val = preliminary_boundaries[range_idx + 1]
            range_size = range_end_val - range_start_val

            new_boundary = range_start_val + (range_size * fraction_into_range)
            # Convert to integer if needed (use round for better distribution)
            if is_integer_field:
                new_boundary = round(new_boundary)
            new_boundaries.append(new_boundary)

        new_boundaries.append(max_value)  # End boundary stays

        # Build optimized ranges
        optimized_ranges = []
        for i in range(partition_count):
            end_val = new_boundaries[i + 1]
            # For integer fields, extend last partition's end by +1 to include max value
            # (since range filter uses lt=end, not lte=end)
            if is_integer_field and i == partition_count - 1:
                end_val = end_val + 1

            optimized_ranges.append(
                {
                    "partition_id": i,
                    "start": new_boundaries[i],
                    "end": end_val,
                    "estimated_count": target_per_partition,
                }
            )

    else:
        raise ValueError(f"Cannot partition by field type: {field_type}")

    # Step 4: Verify actual counts for optimized ranges
    print("  Step 4: Verifying actual counts for optimized ranges...")
    for r in optimized_ranges:
        try:
            result = client.data_modeling.instances.aggregate(
                view=view_id,
                aggregates=[{"count": {"property": "externalId"}}],
                filter=And(
                    Equals(["node", "space"], instance_space),
                    RangeFilter(
                        property=[view_id.space, f"{view_id.external_id}/{view_id.version}", partition_field],
                        gte=r["start"],
                        lt=r["end"],
                    ),
                ),
                instance_type="node",
            )
            actual_count = result[0].value if result else 0
            r["actual_count"] = actual_count
            deviation = (
                abs(actual_count - target_per_partition) / target_per_partition if target_per_partition > 0 else 0
            )
            print(f"    Partition {r['partition_id']}: {actual_count:,} instances ({deviation * 100:.1f}% deviation)")
        except Exception as e:
            print(f"    Partition {r['partition_id']}: Error verifying - {e}")
            r["actual_count"] = r["estimated_count"]  # Fallback to estimate

    # Print final summary
    print("  Final optimized partition ranges:")
    max_deviation = (
        max(
            abs(r.get("actual_count", target_per_partition) - target_per_partition) / target_per_partition
            for r in optimized_ranges
        )
        if optimized_ranges
        else 0
    )
    print(f"    Max deviation from target: {max_deviation * 100:.1f}%")
    for r in optimized_ranges:
        print(f"    Partition {r['partition_id']}: {r.get('actual_count', 0):,} instances ({r['start']} to {r['end']})")

    return optimized_ranges


def trigger_validation_function(
    client: CogniteClient,
    validation_function_id: str,
    partition_data: dict,
    wait: bool = False,
) -> dict:
    """Trigger a single validation function for a partition."""
    print(f"  Triggering partition {partition_data['partition_id']}...")

    call = client.functions.call(
        external_id=validation_function_id,
        data=partition_data,
        wait=wait,
    )

    return {
        "partition_id": partition_data["partition_id"],
        "call_id": call.id,
        "status": call.status if wait else "scheduled",
    }


def handle_orchestrator(
    data: dict[str, Any], client: CogniteClient, secrets: dict[str, str] | None = None
) -> dict[str, Any]:
    """
    Orchestrator function handler with two modes:

    INITIAL MODE (no orchestration_id):
        - Create new orchestration_id
        - Discover data distribution
        - Calculate partition ranges
        - Create self-managing scheduled workflow (every 2 min)
        - Trigger all partition functions
        - Save initial state
        - Return immediately

    MONITOR MODE (orchestration_id provided):
        - Load orchestration state
        - Check all partition statuses
        - Retrigger failed/incomplete partitions (max retries)
        - Update state with progress
        - If all complete: delete workflow, cleanup state
        - If not complete: save updated state
        - Return status

    Parameters (from data):
        orchestration_id: str (optional) - If provided, run in monitor mode

        # INITIAL MODE parameters:
        partition_count: int - Number of partitions (default: 10)
        partition_field: str - Field to partition by (e.g., "scheduledStartTime")
        max_retries: int - Max retries for errors (default: 3)
        max_partial_retries: int - Max retries for timeouts/partial (default: 100)

        # View/Model config
        datamodel_space: str
        datamodel_external_id: str
        datamodel_version: str
        view_space: str (optional, defaults to datamodel_space)
        view_external_id: str (optional, defaults to datamodel_external_id)
        view_version: str (optional, defaults to datamodel_version)
        instance_space: str

        # Validation config (passed through to validation functions)
        shacl_file_external_id: str
        rule_set_id: str
        rule_set_version: str
        stream_id: str
        records_space: str (optional, default: "dataQuality")
        records_container: str (optional, default: "DataQualityValidationRecord")
        data_quality_space: str (optional, default: "dataQuality") - Space for state containers
        auto_load_depth: int (optional)
    """
    start_time = time.time()

    orchestration_id = data.get("orchestration_id")

    # Get data quality space from configuration (default: "dataQuality")
    data_quality_space = data.get("data_quality_space", "dataQuality")

    if orchestration_id:
        # ==============================================================================
        # MONITOR MODE: Check status, retrigger incomplete, update state
        # ==============================================================================
        print("=" * 80)
        print("MONITOR MODE: Checking orchestration status")
        print("=" * 80)
        print(f"Orchestration ID: {orchestration_id}")

        # Load state
        state = load_orchestration_state(client, orchestration_id, data_quality_space)
        if not state:
            # No state found - orchestration already completed and cleaned up
            print(f"No state found for {orchestration_id} - already completed and cleaned up")

            # Only delete the schedule if it is still tracking THIS orchestration_id.
            # A fresh initial run may have already replaced it with a new orchestration.
            view_external_id = data.get("view_external_id")
            if view_external_id:
                try:
                    schedule_name = f"Monitor {view_external_id} validation"
                    schedules = client.functions.schedules.list(limit=-1)
                    for schedule in schedules:
                        if schedule.name == schedule_name:
                            schedule_orch_id = (schedule.get_input_data() or {}).get("orchestration_id")
                            if schedule_orch_id == orchestration_id:
                                delete_function_schedule(client, view_external_id)
                            else:
                                print(
                                    f"  Schedule now tracks {schedule_orch_id}"
                                    f" (not {orchestration_id}), skipping delete"
                                )
                            break
                except Exception as e:
                    print(f"  Could not check/delete schedule: {e}")

            print("=" * 80)
            return {
                "status": "completed",
                "message": "Orchestration already completed and cleaned up",
                "orchestration_id": orchestration_id,
                "runtime": time.time() - start_time,
            }

        print(f"View: {state.get('view_external_id')}")
        print(f"Started: {state.get('created_time')}")
        print(f"Status: {state.get('status')}")
        print(f"Total partitions: {state.get('partition_count')}")
        print("=" * 80)

        # Check if in timestamp collection phase
        collection_phase = state.get("timestamp_collection_phase")
        if collection_phase in ("collecting", "checkpoint"):
            print("⚠ Orchestration is in timestamp collection phase")
            print(f"  Phase: {collection_phase}")
            print(f"  Timestamps collected so far: {state.get('timestamp_collection_count', 0):,}")

            if collection_phase == "checkpoint":
                # Controlled shutdown — function saved state cleanly and exited.
                # 1. Update phase to "collecting" immediately so concurrent monitors don't re-trigger.
                # 2. Run initial mode in this process via _resume_orchestration_id (no chaining).
                print("  Controlled shutdown detected — resuming collection in-process")
                try:
                    state["timestamp_collection_phase"] = "collecting"
                    state["last_updated"] = format_timestamp_for_dms(datetime.now(timezone.utc))
                    save_orchestration_state(client, state, data_quality_space)
                    print("  ✓ Updated phase: checkpoint → collecting")
                except Exception as phase_err:
                    print(f"  ⚠ Could not update phase to collecting: {phase_err}")
                resume_data = {k: v for k, v in data.items() if k != "orchestration_id"}
                resume_data["_resume_orchestration_id"] = orchestration_id
                return handle_orchestrator(resume_data, client, secrets)

            # phase == "collecting": function is (or was) actively running
            # Determine how stale the state is
            last_updated = state.get("last_updated")
            age_seconds = 9999
            if last_updated:
                if hasattr(last_updated, "timestamp"):
                    age_seconds = time.time() - last_updated.timestamp()
                elif isinstance(last_updated, str):
                    try:
                        import datetime as _dt

                        lu_dt = _dt.datetime.fromisoformat(last_updated.replace("Z", "+00:00"))
                        age_seconds = time.time() - lu_dt.timestamp()
                    except Exception:
                        pass

            if age_seconds < 900:  # 15 minutes: initial function likely still running or just crashed
                print(f"  Last updated {age_seconds:.0f}s ago — collection likely in progress")
                print("  Monitor mode: no action needed")
                return {
                    "status": "collecting_timestamps",
                    "orchestration_id": orchestration_id,
                    "timestamps_collected": state.get("timestamp_collection_count", 0),
                    "runtime": time.time() - start_time,
                    "message": "Timestamp collection in progress",
                }
            else:
                # State is stale — initial function must have crashed.
                # Delete the stale state so the next invocation of this (old) schedule
                # sees no state and cleans up gracefully without touching the new schedule.
                # Then trigger a completely fresh initial run (new orchestration_id).
                print(f"  State is stale ({age_seconds:.0f}s since last update) — triggering fresh collection")
                try:
                    delete_orchestration_state(client, orchestration_id, data_quality_space)
                    print("  ✓ Deleted stale state")
                except Exception as del_err:
                    print(f"  ⚠ Could not delete stale state: {del_err}")

                function_ext_id = data.get("function_external_id", "data-quality-validation")
                # Strip orchestration_id so the triggered function runs in initial mode
                initial_data = {k: v for k, v in data.items() if k != "orchestration_id"}
                try:
                    client.functions.call(external_id=function_ext_id, data=initial_data, wait=False)
                    print("  ✓ Fresh initial collection triggered (previous progress will be lost)")
                except Exception as retrigger_err:
                    print(f"  ✗ Failed to trigger fresh initial collection: {retrigger_err}")
                return {
                    "status": "retriggering_fresh_collection",
                    "orchestration_id": orchestration_id,
                    "timestamps_collected": state.get("timestamp_collection_count", 0),
                    "stale_seconds": age_seconds,
                    "runtime": time.time() - start_time,
                    "message": "Stale collection detected, triggered fresh initial run",
                }

        # Get retry limits from workflow config (data), not from state (not stored in container)
        max_retries = data.get("max_retries", 3)
        max_partial_retries = data.get("max_partial_retries", 100)

        # Check partition statuses
        print("Checking partition statuses...")
        completed, failed, in_progress, completed_dict, failed_dict, total_validated, total_violations = (
            check_partition_statuses(
                client, state, max_retries, data.get("function_external_id", "data-quality-validation")
            )
        )

        print("Status summary:")
        print(f"  Completed: {len(completed)}")
        print(f"  Failed: {len(failed)}")
        print(f"  In progress: {len(in_progress)}")

        # Retrigger incomplete partitions
        partition_call_ids = state.get("partition_call_ids", {})
        partition_ranges_dict = state.get("partition_ranges", {})
        partition_ranges = list(partition_ranges_dict.values())  # Convert dict back to list
        partition_retry_count = state.get("partition_retry_count", {})

        # Note: completed_dict and failed_dict come from check_partition_statuses() above

        # Build base config for partitions (from data, not state)
        partition_base_config = {
            "validation_type": "partitioned",
            "shacl_file_external_id": data["shacl_file_external_id"],
            "datamodel_space": data.get("datamodel_space"),
            "datamodel_external_id": data.get("datamodel_external_id"),
            "datamodel_version": data.get("datamodel_version"),
            "view_space": state.get("view_space"),
            "view_external_id": state.get("view_external_id"),
            "view_version": data.get("view_version", data.get("datamodel_version")),
            "instance_space": state.get("instance_space"),
            "rule_set_id": state.get("rule_set_id"),
            "rule_set_version": data.get("rule_set_version", "1.0"),
            "stream_id": data["stream_id"],
            "records_space": data.get("records_space", "dataQuality"),
            "records_container": data.get("records_container", "DataQualityValidationRecord"),
            "data_quality_space": data_quality_space,  # Pass to partitioned functions
            "auto_load_depth": data.get("auto_load_depth", 2),
            "chunk_size": data.get("chunk_size", 200),  # Instances per validation chunk
            "orchestration_id": orchestration_id,
            "job_run_id": orchestration_id,  # Use orchestration_id as job_run_id (shared across all partitions)
        }

        # Iterate over a copy to avoid modifying list during iteration
        for partition_id in list(in_progress):
            old_call_id = partition_call_ids.get(partition_id)
            if not old_call_id:
                continue

            try:
                call = client.functions.calls.retrieve(
                    call_id=int(old_call_id),
                    function_external_id=data.get("function_external_id", "data-quality-validation"),
                )

                # Handle case where function was redeployed (call_id no longer valid)
                if call is None:
                    print(
                        f"  Partition {partition_id}: Call not found"
                        " (function may have been redeployed), will retrigger"
                    )
                    needs_retrigger = True
                    failure_type = "error"  # Treat as error for retry logic
                else:
                    print(f"  Partition {partition_id}: {call.status}")

                    # Check if needs retriggering and determine failure type
                    needs_retrigger = False
                    failure_type = None

                    if call.status in ["Failed", "Timeout"]:
                        needs_retrigger = True
                        failure_type = "error"
                    elif call.status == "Completed":
                        response = call.get_response() if hasattr(call, "get_response") else call.response
                        if response:
                            response_status = response.get("status")
                            if response_status == "error":
                                needs_retrigger = True
                                failure_type = "error"
                            elif response_status == "partial":
                                needs_retrigger = True
                                failure_type = "partial"

                if needs_retrigger:
                    retry_count = partition_retry_count.get(partition_id, 0)

                    # Check appropriate retry limit based on failure type
                    if failure_type == "partial":
                        limit = max_partial_retries
                        limit_name = "partial retries"
                    else:  # error
                        limit = max_retries
                        limit_name = "error retries"

                    if retry_count >= limit:
                        print(f"    Max {limit_name} ({limit}) exceeded, marking as failed")
                        failed_dict[partition_id] = True
                        failed.append(partition_id)
                        in_progress.remove(partition_id)
                    else:
                        # Find partition range config
                        # Note: partition_id is string (from dict keys), pr["partition_id"] is int
                        partition_range = next(
                            (pr for pr in partition_ranges if pr["partition_id"] == int(partition_id)), None
                        )
                        if partition_range:
                            # Rebuild partition_data
                            partition_data = {
                                **partition_base_config,
                                "partition_id": partition_id,
                                "partition_range": {
                                    "min_value": partition_range.get("start"),
                                    "max_value": partition_range.get("end"),
                                },
                                "partition_field": state.get("partition_field"),
                            }

                            # Retrigger
                            print(f"    Retriggering (attempt {retry_count + 1}/{limit}) [{failure_type}]...")
                            call_result = trigger_validation_function(
                                client,
                                data.get("function_external_id", "data-quality-validation"),
                                partition_data,
                                wait=False,
                            )
                            partition_call_ids[partition_id] = call_result["call_id"]
                            partition_retry_count[partition_id] = retry_count + 1
                            print(f"    Retriggered: {call_result['call_id']}")
                        else:
                            print(f"    Warning: No partition range found for {partition_id}")

            except Exception as e:
                print(f"  Partition {partition_id}: Error checking/retriggering - {e}")

        # Update state
        state["partition_call_ids"] = partition_call_ids
        state["partition_retry_count"] = partition_retry_count
        state["completed_partitions"] = completed_dict
        state["failed_partitions"] = failed_dict
        state["validated_instances"] = total_validated
        state["total_violations"] = total_violations
        state["last_updated"] = format_timestamp_for_dms(datetime.now(timezone.utc))

        # Check if complete
        if len(in_progress) == 0:
            if len(failed) == 0:
                state["status"] = "completed"
                print("All partitions completed successfully!")
            else:
                state["status"] = "failed"
                print(f"Orchestration failed: {len(failed)} partitions failed after max retries")

            # Delete function schedule
            view_external_id = state.get("view_external_id")
            if view_external_id:
                delete_function_schedule(client, view_external_id)

            # Save final state (keep for reference)
            save_orchestration_state(client, state, data_quality_space)

            print("=" * 80)
            print("ORCHESTRATION COMPLETE")
            print("=" * 80)

            return {
                "status": state["status"],
                "orchestration_id": orchestration_id,
                "completed": len(completed),
                "failed": len(failed),
                "total": state.get("partition_count"),
                "runtime": time.time() - start_time,
            }
        else:
            # Still in progress, save and continue
            state["status"] = "in_progress"
            save_orchestration_state(client, state, data_quality_space)

            print(f"Orchestration in progress ({len(in_progress)} partitions remaining)")
            print("=" * 80)

            return {
                "status": "in_progress",
                "orchestration_id": orchestration_id,
                "completed": len(completed),
                "in_progress": len(in_progress),
                "failed": len(failed),
                "total": state.get("partition_count"),
                "runtime": time.time() - start_time,
            }

    else:
        # ==============================================================================
        # INITIAL MODE: Setup, trigger, create workflow
        # ==============================================================================
        print("=" * 80)
        print("INITIAL MODE: Setting up orchestration")
        print("=" * 80)

        # Get parameters
        partition_count = data.get("partition_count", 10)
        if partition_count <= 0:
            return {
                "status": "failed",
                "error": f"Invalid partition_count: {partition_count}. Must be > 0.",
                "runtime": time.time() - start_time,
            }

        partition_field = data.get("partition_field", "lastUpdatedTime")
        max_retries = data.get("max_retries", 3)
        max_partial_retries = data.get("max_partial_retries", 100)

        # View configuration
        datamodel_space = data["datamodel_space"]
        datamodel_external_id = data["datamodel_external_id"]
        datamodel_version = data["datamodel_version"]
        view_space = data.get("view_space", datamodel_space)
        view_external_id = data.get("view_external_id", datamodel_external_id)
        view_version = data.get("view_version", datamodel_version)
        instance_space = data["instance_space"]

        # Generate orchestration_id (also used as job_run_id for all partitions).
        # _resume_orchestration_id is set when monitor resumes a checkpoint in-process.
        _resume_orch_id = data.get("_resume_orchestration_id")
        if _resume_orch_id:
            orchestration_id = _resume_orch_id
            print(f"In-process checkpoint resume: {orchestration_id}")
        else:
            timestamp = int(time.time())
            orchestration_id = f"job_{view_external_id}_{timestamp}"

        print(f"Orchestration ID: {orchestration_id}")
        print(f"Job Run ID (orchestration_id): {orchestration_id} (shared across all partitions)")
        print(f"Target function: {data.get('function_external_id', 'data-quality-validation')}")
        print(f"Partition count: {partition_count}")
        print(f"Partition field: {partition_field}")
        print(f"View: {view_space}.{view_external_id} v{view_version}")
        print(f"Instance space: {instance_space}")
        print("=" * 80)

        view_id = ViewId(view_space, view_external_id, view_version)

        # Check if we're resuming timestamp collection
        print("Checking for partial timestamp collection state...")
        collection_state = load_timestamp_collection_state(client, orchestration_id, data_quality_space)

        if collection_state:
            # Check if collection is CURRENTLY RUNNING (not just partial).
            # Skip this check for in-process checkpoint resumes — no concurrency risk.
            last_updated_ms = collection_state.get("last_updated_ms", 0)
            now_ms = int(time.time() * 1000)
            age_seconds = (now_ms - last_updated_ms) / 1000

            if age_seconds < 600 and not _resume_orch_id:  # 10 minutes (function timeout)
                print(f"⚠ Timestamp collection already in progress (last updated {age_seconds:.0f}s ago)")
                print("  Exiting to prevent concurrent execution")
                print("  If this is stuck, check orchestration state and delete if needed")
                return {
                    "status": "already_running",
                    "orchestration_id": orchestration_id,
                    "age_seconds": age_seconds,
                    "message": "Timestamp collection already in progress, skipping to prevent concurrency",
                }

            if _resume_orch_id:
                print("✓ Resuming from checkpoint (in-process)...")
            else:
                # State is old (> 10 min) - previous run timed out, safe to resume.
                print("✓ Found stale partial state, resuming from CogniteFile...")
                print(f"  Last updated: {age_seconds:.0f}s ago (previous run timed out)")
            print(f"  Already collected: {collection_state['count']:,} timestamps")
            print("  Resuming from cursor...")

            all_timestamps = collection_state["timestamps"]
            current_cursor = collection_state["cursor"]

            # Skip schedule creation — it already exists from the previous run
            resuming_collection = True
        else:
            print("  No partial state found, starting fresh")
            resuming_collection = False

        # Step 1: Create function schedule FIRST (so timeout can auto-retry)
        # Only create if NOT resuming (schedule already exists if resuming)
        schedule_id = None  # defined here in case we skip creation (resuming)
        if not resuming_collection:
            print("Step 1: Creating function schedule for monitoring...")
            print("  (Creating schedule first so timestamp collection timeout can auto-retry)")
            schedule_id = create_function_schedule(
                client,
                view_external_id,
                orchestration_id,
                data,  # Pass all original config
                secrets,
            )
            print(f"  ✓ Function schedule created: {schedule_id}")
            print("  Schedule will retry this function every 2 minutes")

        # Step 2: Collect timestamps (if partitioning by lastUpdatedTime)
        # Optimization: Collect all values in one pass to avoid aggregate queries
        if partition_field == "lastUpdatedTime":
            print("Step 2: Collecting timestamps for partition distribution...")
            print("  Partitioning by lastUpdatedTime - collecting distribution in one pass...")
            print("  Using sync with twoPhase=True for efficient metadata collection...")

            # Import Query classes for sync
            from cognite.client.data_classes.data_modeling.query import (
                NodeResultSetExpression,
                Query,
                Select,
            )
            from cognite.client.data_classes.filters import And, HasData

            # Collect all lastUpdatedTime values using sync with twoPhase=True
            iteration_start = time.time()

            # Initialize collection state (or resume from loaded state)
            if not resuming_collection:
                all_timestamps = []
                chunk_count = 0
                current_cursor = None
                # Save initial "collecting" state immediately so the monitor schedule
                # (which fires every 2 min) can see we're collecting and not delete the schedule
                print("  Saving initial collecting state...")
                save_timestamp_collection_state(client, orchestration_id, None, [], data_quality_space)
                print("  ✓ Initial state saved - monitor will wait for collection to complete")
            else:
                # all_timestamps and current_cursor already set from loaded state
                chunk_count = 0  # Reset chunk count for this session
                print(f"  Resuming with {len(all_timestamps):,} timestamps already collected")

            MAX_COLLECTION_TIME = 480  # 8 minutes (2-min safety margin before 10-min timeout)
            STATE_SAVE_INTERVAL = 120  # Save state every 2 minutes
            last_state_save = time.time()

            # Sync limit sequence: each level is tried twice before stepping down.
            # The index persists across chunks so we don't restart at 2000 after a timeout.
            SYNC_LIMITS = [2000, 1000, 500, 200, 100, 10]
            sync_limit_idx = 0

            while True:
                # Check timeout BEFORE sync query (critical!)
                elapsed = time.time() - iteration_start
                if elapsed > MAX_COLLECTION_TIME:
                    print(f"⚠ Approaching timeout ({elapsed:.1f}s), saving phase state and exiting...")
                    print("  Saving timestamps + cursor to CogniteFile for resumable collection...")
                    try:
                        save_timestamp_collection_state(
                            client,
                            orchestration_id,
                            current_cursor,
                            all_timestamps,
                            data_quality_space,
                            phase="checkpoint",
                        )
                        print(f"  ✓ Phase state saved ({len(all_timestamps):,} timestamps counted)")
                    except Exception as save_error:
                        print(f"  ⚠ Warning: Phase state save failed (non-fatal): {save_error}")
                    return {
                        "status": "partial_timestamp_collection",
                        "orchestration_id": orchestration_id,
                        "timestamps_collected": len(all_timestamps),
                        "elapsed_seconds": elapsed,
                        "runtime": time.time() - start_time,
                        "message": "Approaching timeout, monitor will retrigger fresh collection",
                    }

                chunk_count += 1

                result = None
                retries_at_limit = 0  # attempts at the current limit level (resets per chunk)

                while True:
                    sync_limit = SYNC_LIMITS[sync_limit_idx]
                    try:
                        # Build sync query with two_phase mode - filter to only nodes with data in view
                        # Select(sources=[]) fetches no view properties — only system metadata
                        # (space, externalId, lastUpdatedTime) which is all we need here.
                        query = Query(
                            with_={
                                "nodes": NodeResultSetExpression(
                                    filter=And(Equals(["node", "space"], instance_space), HasData(views=[view_id])),
                                    limit=sync_limit,
                                    sync_mode="two_phase",
                                )
                            },
                            select={"nodes": Select(sources=[])},
                            cursors={"nodes": current_cursor} if current_cursor else None,
                        )

                        # Use sync to get only metadata (including lastUpdatedTime)
                        result = client.data_modeling.instances.sync(query=query)
                        break  # Success, exit retry loop

                    except Exception as e:
                        error_msg = str(e).lower()
                        if "timeout" in error_msg or "timed out" in error_msg:
                            retries_at_limit += 1
                            if retries_at_limit < 2:
                                # Second attempt at the same limit
                                print(
                                    f"    Sync timeout (attempt {retries_at_limit}/2),"
                                    f" retrying with same limit={sync_limit} in 30s..."
                                )
                                time.sleep(30)
                            else:
                                # Both attempts at this limit failed — step down
                                retries_at_limit = 0
                                sync_limit_idx += 1
                                if sync_limit_idx >= len(SYNC_LIMITS):
                                    print("    All sync limits exhausted, saving progress for retry...")
                                    try:
                                        save_timestamp_collection_state(
                                            client,
                                            orchestration_id,
                                            current_cursor,
                                            all_timestamps,
                                            data_quality_space,
                                        )
                                        print(f"    ✓ Saved {len(all_timestamps):,} timestamps with cursor")
                                    except Exception as save_err:
                                        print(f"    ⚠ Warning: Could not save progress: {save_err}")
                                    return {
                                        "status": "sync_limits_exhausted",
                                        "orchestration_id": orchestration_id,
                                        "timestamps_collected": len(all_timestamps),
                                        "runtime": time.time() - start_time,
                                        "message": (
                                            "All sync limits exhausted (including 10),"
                                            " saved progress, monitor will retry"
                                        ),
                                    }
                                print(
                                    f"    Sync timeout (2/2), reducing limit: {sync_limit}"
                                    f" → {SYNC_LIMITS[sync_limit_idx]}, waiting 30s..."
                                )
                                time.sleep(30)
                        else:
                            # Not a timeout error, raise immediately
                            raise

                if result is None:
                    raise RuntimeError("Sync query failed - no result returned")

                # Extract lastUpdatedTime from nodes
                chunk_size = 0
                if hasattr(result, "data") and "nodes" in result.data:
                    for node in result.data["nodes"]:
                        all_timestamps.append(node.last_updated_time)
                        chunk_size += 1

                # Save lightweight progress state every 2 minutes (phase + count only,
                # no timestamp list — DMS JSON property limit is 40KB)
                if chunk_size > 0 and (time.time() - last_state_save) >= STATE_SAVE_INTERVAL:
                    try:
                        save_timestamp_collection_state(
                            client, orchestration_id, current_cursor, all_timestamps, data_quality_space
                        )
                        last_state_save = time.time()
                    except Exception as save_error:
                        # Non-fatal: state saves are only for monitoring visibility.
                        # Timestamps are kept in memory — collection continues regardless.
                        print(f"  ⚠ Warning: State save failed (continuing): {save_error}")

                # Progress logging every 10 chunks
                if chunk_count % 10 == 0:
                    elapsed = time.time() - iteration_start
                    rate = len(all_timestamps) / elapsed if elapsed > 0 else 0
                    print(
                        f"    Progress: {len(all_timestamps):,} timestamps collected"
                        f" ({rate:.0f}/sec, {chunk_count} chunks, {elapsed:.1f}s)"
                    )

                # Break if no data received (backfill phase complete)
                if chunk_size == 0:
                    print("    Backfill complete: no more data")
                    break

                # Get cursor for next iteration
                new_cursor = None
                if hasattr(result, "cursors"):
                    if hasattr(result.cursors, "__dict__"):
                        new_cursor = result.cursors.__dict__.get("nodes")
                    elif isinstance(result.cursors, dict):
                        new_cursor = result.cursors.get("nodes")

                # Check if more data exists
                has_more = new_cursor is not None and new_cursor != current_cursor
                current_cursor = new_cursor

                if not has_more:
                    break

            iteration_duration = time.time() - iteration_start

            if not all_timestamps:
                return {
                    "status": "failed",
                    "error": "No instances found in the specified space",
                    "runtime": time.time() - start_time,
                }

            # Sort timestamps
            all_timestamps.sort()
            total_count = len(all_timestamps)
            min_timestamp_ms = all_timestamps[0]
            max_timestamp_ms = all_timestamps[-1]

            # Mark timestamp collection complete
            print("✓ Timestamp collection complete!")
            print(f"  Total: {total_count:,} timestamps")
            print(f"  Time: {iteration_duration:.1f}s")
            try:
                mark_timestamp_collection_complete(client, orchestration_id, all_timestamps, data_quality_space)
            except Exception as complete_error:
                print(f"  ⚠ Warning: Failed to mark collection complete: {complete_error}")
                # Continue anyway - collection is done, this is just state management

            # Convert to datetime strings
            min_value = format_timestamp_for_dms(datetime.fromtimestamp(min_timestamp_ms / 1000.0, tz=timezone.utc))
            max_value = format_timestamp_for_dms(datetime.fromtimestamp(max_timestamp_ms / 1000.0, tz=timezone.utc))

            instances_per_sec = total_count / iteration_duration if iteration_duration > 0 else 0
            print(f"  ✓ Collected {total_count:,} timestamps from {chunk_count} chunks")
            print(f"    Iteration time: {iteration_duration:.2f}s ({instances_per_sec:.0f} instances/sec)")
            print(f"    Range: {min_value} to {max_value}")

            distribution = {
                "total_count": total_count,
                "min_value": min_value,
                "max_value": max_value,
                "field_type": "datetime",
            }

            # Calculate partition ranges from actual distribution
            print("Calculating partition ranges from actual distribution...")
            target_per_partition = total_count // partition_count
            partition_ranges = []

            for i in range(partition_count):
                start_idx = i * target_per_partition
                if i == partition_count - 1:
                    # Last partition gets remainder
                    end_idx = total_count - 1
                else:
                    end_idx = (i + 1) * target_per_partition - 1

                start_timestamp_ms = all_timestamps[start_idx]
                end_timestamp_ms = all_timestamps[end_idx]

                partition_ranges.append(
                    {
                        "partition_id": i,
                        "start": format_timestamp_for_dms(
                            datetime.fromtimestamp(start_timestamp_ms / 1000.0, tz=timezone.utc)
                        ),
                        "end": format_timestamp_for_dms(
                            datetime.fromtimestamp(end_timestamp_ms / 1000.0, tz=timezone.utc)
                        ),
                        "actual_count": end_idx - start_idx + 1,
                    }
                )

            # Sync trigger uses max value
            sync_cutoff_timestamp = max_value
            max_partition_value = partition_ranges[-1]["end"]  # For result output

            print(f"  ✓ Created {partition_count} partitions from actual distribution")
            for r in partition_ranges:
                print(
                    f"    Partition {r['partition_id']}: {r['actual_count']:,} instances ({r['start']} to {r['end']})"
                )

        else:
            # Normal flow for non-lastUpdatedTime fields
            # Step 1: Discover data distribution
            print("Step 1: Discovering data distribution...")

            try:
                distribution = discover_data_distribution(
                    client,
                    instance_space,
                    view_id,
                    partition_field,
                )
            except Exception as e:
                return {
                    "status": "failed",
                    "error": f"Failed to discover data: {e!s}",
                    "runtime": time.time() - start_time,
                }

            # Step 2: Calculate partition ranges
            print("Step 2: Calculating partition ranges...")
            try:
                partition_ranges = calculate_partition_ranges(
                    client,
                    distribution,
                    partition_count,
                    instance_space,
                    view_id,
                    partition_field,
                )
            except Exception as e:
                return {
                    "status": "failed",
                    "error": f"Failed to calculate partitions: {e!s}",
                    "distribution": distribution,
                    "runtime": time.time() - start_time,
                }

            # Step 2.5: Determine sync trigger cutoff
            max_partition_value = partition_ranges[-1]["end"]

            if distribution["field_type"] == "datetime":
                # Partition field is a timestamp - use its max value
                sync_cutoff_timestamp = max_partition_value
                print(f"Sync trigger cutoff: {sync_cutoff_timestamp} (from partition field)")
            else:
                # Partition field is not a timestamp - query for max lastUpdatedTime
                # TODO: Replace with aggregate once DMS supports max(lastUpdatedTime)
                # Currently aggregate throws "Expected property to be of a numerical type: lastUpdatedTime"
                # This is a bug - lastUpdatedTime is stored as int64 (milliseconds) and should support min/max
                # Once fixed, use: client.data_modeling.instances.aggregate(
                #     view=view_id, aggregates=[{"max": {"property": "lastUpdatedTime"}}], ...)

                print("Partition field is non-timestamp, querying for max lastUpdatedTime...")
                print("  Using sync with twoPhase=True for efficient metadata collection...")

                # Import Query classes for sync
                from cognite.client.data_classes.data_modeling.query import (
                    NodeResultSetExpression,
                    Query,
                    Select,
                )
                from cognite.client.data_classes.filters import And, HasData

                iteration_start = time.time()
                max_last_updated = 0
                chunk_count = 0
                instance_count = 0
                current_cursor = None

                while True:
                    chunk_count += 1

                    # Build sync query with two_phase mode - filter to only nodes with data in view
                    # Select(sources=[]) fetches no view properties — only system metadata
                    # (space, externalId, lastUpdatedTime) which is all we need here.
                    query = Query(
                        with_={
                            "nodes": NodeResultSetExpression(
                                filter=And(Equals(["node", "space"], instance_space), HasData(views=[view_id])),
                                limit=1000,
                                sync_mode="two_phase",
                            )
                        },
                        select={"nodes": Select(sources=[])},
                        cursors={"nodes": current_cursor} if current_cursor else None,
                    )

                    # Use sync to get only metadata (including lastUpdatedTime)
                    result = client.data_modeling.instances.sync(query=query)

                    # Extract lastUpdatedTime from nodes
                    chunk_size = 0
                    if hasattr(result, "data") and "nodes" in result.data:
                        for node in result.data["nodes"]:
                            instance_count += 1
                            chunk_size += 1
                            if node.last_updated_time > max_last_updated:
                                max_last_updated = node.last_updated_time

                    # Progress logging every 10 chunks
                    if chunk_count % 10 == 0:
                        elapsed = time.time() - iteration_start
                        rate = instance_count / elapsed if elapsed > 0 else 0
                        print(
                            f"    Progress: {instance_count:,} instances processed"
                            f" ({rate:.0f}/sec, {chunk_count} chunks, {elapsed:.1f}s)"
                        )

                    # Break if no data received (backfill phase complete)
                    if chunk_size == 0:
                        print("    Backfill complete: no more data")
                        break

                    # Get cursor for next iteration
                    new_cursor = None
                    if hasattr(result, "cursors"):
                        if hasattr(result.cursors, "__dict__"):
                            new_cursor = result.cursors.__dict__.get("nodes")
                        elif isinstance(result.cursors, dict):
                            new_cursor = result.cursors.get("nodes")

                    # Check if more data exists
                    has_more = new_cursor is not None and new_cursor != current_cursor
                    current_cursor = new_cursor

                    if not has_more:
                        break

                iteration_duration = time.time() - iteration_start

                if max_last_updated > 0:
                    instances_per_sec = instance_count / iteration_duration if iteration_duration > 0 else 0
                    sync_cutoff_timestamp = format_timestamp_for_dms(
                        datetime.fromtimestamp(max_last_updated / 1000.0, tz=timezone.utc)
                    )
                    print(f"  ✓ Max lastUpdatedTime: {sync_cutoff_timestamp}")
                    print(
                        f"    Iteration time: {iteration_duration:.2f}s"
                        f" ({instances_per_sec:.0f} instances/sec, {chunk_count} chunks)"
                    )
                else:
                    raise ValueError("No instances found to determine max lastUpdatedTime")

        # Step 3: Create sync trigger
        print("Step 3: Creating sync trigger for ongoing validation...")
        print(f"  Historic cutoff: {sync_cutoff_timestamp}")
        print(f"  Historic partitions by: {partition_field}")

        # Construct sync workflow external_id using same logic as deploy_workflows.py
        default_sync_workflow_id = f"dq-shacl-{view_external_id.lower()}"

        sync_trigger_id = create_sync_trigger(
            client,
            {
                "view_space": view_space,
                "view_external_id": view_external_id,
                "view_version": view_version,
                "instance_space": instance_space,
                "workflow_external_id": data.get("sync_workflow_external_id", default_sync_workflow_id),
                "datamodel_space": datamodel_space,
                "datamodel_external_id": datamodel_external_id,
                "datamodel_version": datamodel_version,
                "sync_batch_size": data.get(
                    "sync_batch_size",
                    1000 if data.get("use_sync_cursor_mode") else 100,
                ),
                "sync_batch_timeout": data.get("sync_batch_timeout", 300),
                "use_sync_cursor_mode": data.get("use_sync_cursor_mode", False),
            },
            sync_cutoff_timestamp,
            orchestration_id,
            orchestration_id,  # Use orchestration_id as job_run_id
            secrets,
        )
        print(f"  ✓ Sync trigger created: {sync_trigger_id}")

        # Step 4: Build partition configs
        print("Step 4: Building partition configs...")

        partition_base_config = {
            "validation_type": "partitioned",
            "shacl_file_external_id": data["shacl_file_external_id"],
            "datamodel_space": datamodel_space,
            "datamodel_external_id": datamodel_external_id,
            "datamodel_version": datamodel_version,
            "view_space": view_space,
            "view_external_id": view_external_id,
            "view_version": view_version,
            "instance_space": instance_space,
            "rule_set_id": data["rule_set_id"],
            "rule_set_version": data.get("rule_set_version", "1.0"),
            "stream_id": data["stream_id"],
            "records_space": data.get("records_space", "dataQuality"),
            "records_container": data.get("records_container", "DataQualityValidationRecord"),
            "data_quality_space": data_quality_space,  # Pass to partitioned functions
            "auto_load_depth": data.get("auto_load_depth", 2),
            "chunk_size": data.get("chunk_size", 200),  # Instances per validation chunk
            "orchestration_id": orchestration_id,  # Include for tracking
            "job_run_id": orchestration_id,  # Use orchestration_id as job_run_id (shared across all partitions)
        }

        # Step 5: Trigger all partitions
        print(f"Step 5: Triggering {partition_count} validation functions...")

        partition_call_ids = {}
        partition_data_map = {}
        partition_retry_count = {}

        for partition_range in partition_ranges:
            partition_id = partition_range["partition_id"]
            partition_data = {
                **partition_base_config,
                "partition_id": partition_id,
                "partition_range": {  # CRITICAL: Pass range boundaries in correct format
                    "min_value": partition_range.get("start"),
                    "max_value": partition_range.get("end"),
                },
                "partition_field": partition_field,
            }
            partition_data_map[partition_id] = partition_data
            partition_retry_count[partition_id] = 0

            try:
                call_result = trigger_validation_function(
                    client,
                    data.get("function_external_id", "data-quality-validation"),
                    partition_data,
                    wait=False,
                )
                partition_call_ids[partition_id] = call_result["call_id"]
                print(f"  Triggered partition {partition_id}: {call_result['call_id']}")
            except Exception as e:
                print(f"  Failed to trigger partition {partition_id}: {e}")

        # Step 6: Save initial state
        print("Step 6: Saving orchestration state...")

        # Only include fields that exist in OrchestrationState container schema
        # Note: max_retries and max_partial_retries are passed via workflow config, not saved to container
        state = {
            "orchestration_id": orchestration_id,
            "view_space": view_space,
            "view_external_id": view_external_id,
            "instance_space": instance_space,
            "rule_set_id": data.get("rule_set_id", f"{view_external_id}_SHACL"),
            "partition_field": partition_field,
            "partition_count": partition_count,
            "partition_ranges": {str(p["partition_id"]): p for p in partition_ranges},  # Convert list to dict
            "partition_call_ids": partition_call_ids,
            "partition_retry_count": partition_retry_count,
            "triggered_partitions": {p: True for p in partition_call_ids.keys()},
            "completed_partitions": {},
            "failed_partitions": {},
            "total_instances": distribution.get("total_count", 0),
            "validated_instances": 0,
            "total_violations": 0,
            "status": "monitoring",
            "created_time": format_timestamp_for_dms(datetime.now(timezone.utc)),
            "last_updated": format_timestamp_for_dms(datetime.now(timezone.utc)),
            "sync_trigger_external_id": sync_trigger_id,
            "historic_cutoff_date": str(sync_cutoff_timestamp),  # Sync trigger timestamp (always a timestamp)
        }

        save_orchestration_state(client, state, data_quality_space)

        elapsed = time.time() - start_time

        print("=" * 80)
        print("ORCHESTRATION STARTED")
        print("=" * 80)
        print(f"Orchestration ID: {orchestration_id}")
        print(f"Partitions triggered: {len(partition_call_ids)}/{partition_count}")
        print(f"Sync trigger: {sync_trigger_id}")
        print(f"Monitor schedule: Monitor {view_external_id} validation")
        print(f"Runtime: {elapsed:.1f}s")
        print("Monitoring schedule will check progress every 2 minutes")
        print("and auto-delete when orchestration completes.")
        print("=" * 80)

        return {
            "status": "started",
            "orchestration_id": orchestration_id,
            "sync_trigger_external_id": sync_trigger_id,
            "monitor_schedule_name": schedule_id,
            "historic_cutoff_date": str(sync_cutoff_timestamp),  # Sync trigger timestamp as string
            "partition_field": partition_field,
            "partition_field_max": str(max_partition_value),  # Max value of partition field
            "partition_count": partition_count,
            "partitions_triggered": len(partition_call_ids),
            "distribution": distribution,
            "runtime": elapsed,
        }
