import pytest
from unittest.mock import patch, MagicMock
from awe.clients.comfyui_client import ComfyUIClient

def test_comfyui_client_init():
    """测试 ComfyUI 客户端初始化"""
    client = ComfyUIClient("localhost", "8188")
    assert client.host == "localhost"
    assert client.port == "8188"
    assert client.status == "INIT"

@patch('websocket.WebSocket')
def test_comfyui_client_connect(mock_websocket):
    """测试 WebSocket 连接"""
    mock_ws = MagicMock()
    mock_websocket.return_value = mock_ws
    
    client = ComfyUIClient("localhost", "8188")
    client.connect()
    
    mock_ws.connect.assert_called_once()
    assert client.ws is not None

def test_comfyui_client_progress():
    """测试进度计算"""
    client = ComfyUIClient("localhost", "8188")
    
    # 测试空节点列表
    client.current_total_node_ids = []
    assert client.get_progress() == 1
    
    # 测试正常进度计算
    client.current_total_node_ids = ["1", "2", "3", "4"]
    client.current_executed_node_ids = ["1", "2"]
    assert client.get_progress() == 50
    
    # 测试进度为0的情况
    client.current_executed_node_ids = []
    assert client.get_progress() == 1  # 应该返回1而不是0 

@patch('awe.clients.comfyui_client.build_result')
@patch('awe.clients.comfyui_client.handle_image_result')
def test_build_result_passes_task_id(mock_handle_image_result, mock_build_result):
    client = ComfyUIClient("localhost", "8188")
    client.current_payload = MagicMock(messageTag='APP', taskId='task-123')
    comfyui_result = {"1": {"images": []}}
    workflow_data = {"outputTpl": {}}
    mock_build_result.return_value = {"ok": True}

    result = client.build_result(comfyui_result, workflow_data)

    assert result == {"ok": True}
    mock_handle_image_result.assert_called_once_with(comfyui_result, client.compress, 'task-123')

@patch('awe.clients.comfyui_client.append_image_path')
def test_handle_progress_event_passes_task_id(mock_append_image_path):
    client = ComfyUIClient("localhost", "8188")
    client.progress_callback = None
    client.current_payload = MagicMock(
        messageTag='WEB',
        taskId='task-abc',
        clientId='client-1'
    )

    message = {
        "type": "executed",
        "data": {
            "output": {
                "images": [{"filename": "a.png"}]
            }
        }
    }

    client.handle_progress_event(message)

    mock_append_image_path.assert_called_once_with(message["data"]["output"]["images"], 'png', client.compress, 'task-abc')


def test_handle_progress_event_marks_executing_as_debounced():
    client = ComfyUIClient("localhost", "8188")
    client.progress_callback = MagicMock()
    client.current_payload = MagicMock(
        messageTag='APP',
        taskId='task-abc',
        clientId='client-1'
    )

    client.handle_progress_event({"type": "executing", "data": {"node": "1"}})

    response_message = client.progress_callback.call_args.args[0]
    assert response_message.debounced is True


def test_handle_error_supports_dict_detail():
    client = ComfyUIClient("localhost", "8188")
    client.error_callback = MagicMock()
    client.current_payload = MagicMock(clientId='client-1', taskId='task-1')

    client.handle_error(
        {"detail": "Error happened #42'"},
        {"message": "bad request", "details": "detail text"},
        {}
    )

    client.error_callback.assert_called_once()
    response_message = client.error_callback.call_args.args[0]
    assert response_message.response["error"]["data"]["node_id"] == "42"