"""
Utility functions for cluster operations.
"""

import json
from typing import Dict


def extract_cluster_ids_from_execution(
    raw_execution_response,
) -> Dict[str, str]:
    """
    Extract cluster IDs from TCMS workflow execution response.

    Args:
        raw_execution_response: Raw execution response from TCMS
                               (can be dict or JSON string)

    Returns:
        Dict mapping cluster_key to cluster_id
        e.g., {"create": "123456", "compare": "789012"}
        Returns empty dict if no cluster IDs found or on error
    """
    try:
        # Parse JSON string if needed
        if isinstance(raw_execution_response, str):
            raw_resp = json.loads(raw_execution_response)
        else:
            raw_resp = raw_execution_response

        if not raw_resp:
            return {}

        clusters = {}

        # Navigate to cluster_id
        graph = raw_resp.get("graph", {})
        completed_tasks_list = graph.get("completedTasks", [])

        for task_item in completed_tasks_list:
            node_name = task_item.get("nodeName", "")
            if "tidbcloud_cluster_" not in node_name:
                continue

            state_output = task_item.get("stateOutput", {})
            testbed_info = state_output.get("testbedInfo", {})
            resource_items = testbed_info.get("resourceItems", [])

            for item in resource_items:
                item_type = item.get("type")
                if item_type == "TIDBCLOUD_CLUSTER":
                    details = item.get("details", {})
                    status = details.get("status", {})
                    cluster_id = status.get("clusterId")

                    if (
                        cluster_id
                        and isinstance(cluster_id, str)
                        and cluster_id.isdigit()
                        and len(cluster_id) > 1
                    ):
                        # Extract key from node name
                        # e.g., "tidbcloud_cluster_create" -> "create"
                        key = node_name.removeprefix("tidbcloud_cluster_")
                        clusters[key] = cluster_id
                        break  # Found cluster for this node

        return clusters

    except Exception as e:
        print(f"Error extracting cluster IDs: {e}")
        return {}