from awe.types.common import ResponseMessage
from awe.util.logging import logger


def format_workflow_resp(resp, output_tpl):
    result_obj = {}
    for key in output_tpl:
        output_config = output_tpl.get(key)
        node_id = output_config.get("nodeId")
        param_key = output_config.get("paramKey")
        node = resp.get(f"{node_id}")
        result = None
        try:
            result = node.get(param_key)
        except Exception as error:
            logger.error(f"get result failed: {error}")
        try:
            if key == "images" or output_config.get("paramType") == "Image":
                if result is None and param_key != "gifs":
                    result = node.get("gifs")
        except Exception as error:
            logger.error(f"get result failed: {error}")
        result_obj[key] = {
            "output": result,
            "title": output_config.get("name"),
            "paramType": output_config.get("paramType"),
        }
    return result_obj


def build_result(comfyui_result, workflow_data):
    output_tpl = workflow_data.get("outputTpl")
    return format_workflow_resp(comfyui_result, output_tpl)


def convert_to_old_data(message: ResponseMessage):
    data = {}
    try:
        message_data = message.response_dict().get("data", {})
        if message_data:
            for key, value in message_data.items():
                data[key] = value["output"]
    except Exception as error:
        logger.error(f"[MQProducer] for compatability, get image output from response error: {error}")
        data = {}
    return data


def convert_to_old_result(message: ResponseMessage):
    results = []
    try:
        response_data = message.response_dict().get("data", {})
        results = response_data.get("images", {}).get("output") or []
    except Exception as error:
        logger.error(f"[MQProducer] for compatability, get image output from response error: {error}")
        results = []
    return results
