import base64
import os
import traceback
from io import BytesIO

from awe.oss import get_download_url, upload_file
from awe.util.comfyui_exception import ComfyUIException
from awe.util.logging import logger
from awe.util.util import (
    generate_filename_with_date,
    get_file_extension,
    get_remote_base64_image,
    upload_file_to_comfyui,
    upload_image_to_comfyui,
    upload_remote_image,
)

MEDIA_INPUT_KEYS = {"image", "audio", "video"}


def build_comfyui_upload_url():
    host = os.getenv("COMFYUI_HOST")
    port = os.getenv("COMFYUI_PORT")
    return f"http://{host}:{port}/upload/image"


def resolve_remote_asset_url(asset_value):
    if asset_value.startswith("http"):
        return asset_value
    download_url = get_download_url(asset_value)
    logger.info(f"resolved asset url: {download_url}")
    return download_url


def build_image_node(param, inputs, node, node_id):
    handled = node.get("handled")
    if handled:
        logger.info(f"[WorkflowUtil] node {node_id} has been handled")
        return

    img_url = param
    class_type = node["class_type"]
    if class_type == "Base64ImageInput":
        inputs["base64_image"] = get_remote_base64_image(img_url)
        return

    logger.info(f"[WorkflowUtil] image url: {img_url}")
    if img_url is None or not isinstance(img_url, str):
        logger.error(f"image is not url, node_id: {node_id}, type: {node['class_type']}")
        return

    new_image_url = resolve_remote_asset_url(img_url)
    resp = upload_image_to_comfyui(build_comfyui_upload_url(), "image.png", new_image_url)
    if resp is not None:
        inputs["image"] = resp.get("name")
        node["handled"] = True


def build_array_node(key, param_key, param_node, param, inputs):
    param_sub_type = param_node.get("paramSubType")
    if not isinstance(param, list):
        if param_sub_type is None and isinstance(param, str):
            inputs[f"{param_key}"] = param
            return
        logger.error(f"param {key} is not list, but paramSubType is Array, skip it.")
        return

    if param_sub_type == "Image":
        inputs[f"{param_key}"] = [get_remote_base64_image(item) for item in param]
    else:
        inputs[f"{param_key}"] = param


def limit_input_value(input_key, inputs, max_value, options=None):
    options = options or {}
    node = options.get("node")
    node_id = options.get("node_id")

    if isinstance(inputs[input_key], str):
        inputs[input_key] = int(inputs[input_key])

    if isinstance(inputs[input_key], (int, float)) and inputs[input_key] > max_value:
        inputs[input_key] = max_value
        logger.warning(f"Limiting {input_key} to {max_value}, {inputs}")
        node_type = node.get("class_type", "Unknown") if node is not None else "Unknown"
        raise ComfyUIException(f"{input_key} 不能大于 {max_value}", node_id, node_type)


def _handle_image_input(inputs, new_inputs, node, node_id):
    try:
        build_image_node(inputs["image"], new_inputs, node, node_id)
    except Exception as error:
        logger.error(f"[WorkflowUtil] handle image failed: {error}")
        raise ComfyUIException("图片有误，请重新上传", node_id, node.get("class_type", "Unknown"))


def _handle_media_input(inputs, new_inputs, node, node_id, media_type):
    try:
        handled = node.get("handled")
        if handled:
            logger.info(f"[WorkflowUtil] node {node_id} has been handled")
            return

        media_url = inputs[media_type]
        class_type = node["class_type"]
        logger.info(f"[WorkflowUtil] media url: {media_url}")
        if media_url is None or not isinstance(media_url, str):
            logger.error(f"media is not url, node_id: {node_id}, type: {class_type}")
            return

        new_media_url = resolve_remote_asset_url(media_url)
        ext = get_file_extension(media_url)
        resp = upload_file_to_comfyui(build_comfyui_upload_url(), f"media.{ext}", new_media_url)
        if resp is not None:
            new_inputs[media_type] = resp.get("name")
            node["handled"] = True
    except Exception as error:
        logger.error(f"[WorkflowUtil] handle {media_type} failed: {error}")
        raise ComfyUIException(f"{media_type} 有误，请重新上传", node_id, node.get("class_type", "Unknown"))


def _process_input_key(key, inputs, new_inputs, node, node_id):
    if key == "image":
        _handle_image_input(inputs, new_inputs, node, node_id)
    elif key in {"audio", "video"}:
        _handle_media_input(inputs, new_inputs, node, node_id, key)


def handle_workflow_inputs(workflow):
    for node_id, node in workflow.items():
        try:
            inputs = node.get("inputs", {})
            new_inputs = {**inputs}

            for key in inputs:
                if key in MEDIA_INPUT_KEYS:
                    _process_input_key(key, inputs, new_inputs, node, node_id)

            node["inputs"] = new_inputs
            workflow[node_id] = node
        except Exception as error:
            logger.error(f"[WorkflowUtil] handle_workflow_inputs failed: {error}")
            logger.error(f"[WorkflowUtil] handle_workflow_inputs failed: {traceback.print_exc()}")
            if isinstance(error, ComfyUIException):
                raise error
            raise ComfyUIException(
                f"handle_workflow_inputs failed: {error}",
                node_id,
                node.get("class_type", "Unknown"),
            )

    return workflow


def build_remote_image_url():
    host = os.getenv("COMFYUI_HOST")
    port = os.getenv("COMFYUI_PORT")
    return f"http://{host}:{port}/view"


def _handle_filename_item(item, file_extension, compress=True, task_id=None):
    if item.get("url") is not None and item.get("url").startswith("http"):
        logger.info(f"item url is {item.get('url')}, return it")
        return item

    path = generate_filename_with_date(get_file_extension(item["filename"]))
    preview_url = build_remote_image_url()
    new_path, file_url = upload_remote_image(preview_url, item, path, file_extension, compress, task_id)
    return {**item, "url": file_url, "path": new_path}


def _handle_id_item(item, compress=True, task_id=None):
    result = {**item}
    urls = {
        "image_url": "png",
        "video_url": "mp4",
        "audio_url": "mp3",
    }

    for url_key, ext in urls.items():
        url = item.get(url_key)
        if url:
            path = generate_filename_with_date(ext)
            new_url, _ = upload_remote_image(url, item, path, ext, compress, task_id)
            result[url_key] = new_url
            if url_key == "image_url":
                result["path"] = new_url

    return result


def _handle_dict_item(item, file_extension, compress=True, task_id=None):
    if item.get("filename"):
        return _handle_filename_item(item, file_extension, compress, task_id)
    if item.get("id"):
        return _handle_id_item(item, compress, task_id)
    logger.error(f"filename not found in item: {item}")
    return item


def _handle_str_item(base64_image):
    if base64_image.startswith("http"):
        return {"url": base64_image, "path": base64_image}

    if base64_image.startswith("data:image"):
        path = generate_filename_with_date("png")
        image_bytes = base64.b64decode(base64_image)
        input_stream = BytesIO(image_bytes)
        file_url = upload_file(input_stream, path)
        return {"url": file_url, "path": path}

    logger.error(f"base64_image {base64_image} not support")
    return base64_image


def append_image_path(items, file_extension="png", compress=True, task_id=None):
    for index in range(len(items)):
        if isinstance(items[index], dict):
            items[index] = _handle_dict_item(items[index], file_extension, compress, task_id)
        elif isinstance(items[index], str):
            items[index] = _handle_str_item(items[index])
        else:
            logger.error(f"item {items[index]} not support")


def handle_image_result(comfyui_result, compress=True, task_id=None):
    try:
        for node_id in comfyui_result:
            node = comfyui_result.get(node_id)
            for key in node:
                item = node.get(key)
                if isinstance(item, dict):
                    item_list = [item]
                    append_image_path(item_list, "png", compress, task_id)
                    node[key] = item_list[0]
                elif isinstance(item, list):
                    append_image_path(item, "png", compress, task_id)
                else:
                    logger.error(f"item {item} not support")
    except Exception as error:
        logger.error(f"[WorkflowUtil] handle_image_result failed: {error}")
        logger.error(f"[WorkflowUtil] handle_image_result failed: {traceback.print_exc()}")
