import pytest
from comfy.client.embedded_comfy_client import Comfy
from comfy.cli_args_types import Configuration
from comfy.execution_context import context_configuration
from comfy.distributed.process_pool_executor import ProcessPoolExecutor
from comfy.distributed.executors import ContextVarExecutor


@pytest.mark.asyncio
async def test_comfy_defaults():
    """Test that default initialization uses ContextVarExecutor."""
    client = Comfy()
    assert isinstance(client._executor, ContextVarExecutor)
    assert client._owns_executor


@pytest.mark.asyncio
async def test_comfy_config_triggers_process_pool():
    """Test that configurations affecting global state trigger ProcessPoolExecutor."""
    # "lowvram" is in MODEL_MANAGEMENT_ARGS
    client = Comfy(configuration={"lowvram": True})
    assert isinstance(client._executor, ProcessPoolExecutor)
    assert client._owns_executor


@pytest.mark.asyncio
async def test_comfy_disable_dynamic_vram_triggers_process_pool():
    """Disabling dynamic VRAM changes model-management behavior."""
    client = Comfy(configuration={"disable_dynamic_vram": True})
    assert isinstance(client._executor, ProcessPoolExecutor)
    assert client._owns_executor


@pytest.mark.asyncio
async def test_comfy_uses_current_context_as_process_pool_baseline():
    """Comparisons should use the current execution context configuration."""
    with context_configuration(Configuration(disable_dynamic_vram=True)):
        matching_client = Comfy(configuration={"disable_dynamic_vram": True})
        assert isinstance(matching_client._executor, ContextVarExecutor)

        differing_client = Comfy(configuration={"disable_dynamic_vram": False})
        assert isinstance(differing_client._executor, ProcessPoolExecutor)


@pytest.mark.asyncio
async def test_comfy_config_unrelated_does_not_trigger():
    """Test that unrelated configuration keys do NOT trigger ProcessPoolExecutor."""
    # "some_random_ui_setting" is likely not in MODEL_MANAGEMENT_ARGS
    client = Comfy(configuration={"some_random_ui_setting": "value"})
    assert isinstance(client._executor, ContextVarExecutor)


@pytest.mark.asyncio
async def test_comfy_explicit_executor_string():
    """Test explicitly requesting an executor by string."""
    client = Comfy(executor="ProcessPoolExecutor")
    assert isinstance(client._executor, ProcessPoolExecutor)
    assert client._owns_executor

    client2 = Comfy(executor="ContextVarExecutor")
    assert isinstance(client2._executor, ContextVarExecutor)
    assert client2._owns_executor


@pytest.mark.asyncio
async def test_comfy_explicit_executor_instance():
    """Test passing an executor instance."""
    executor = ContextVarExecutor(max_workers=1)
    client = Comfy(executor=executor)
    assert client._executor is executor
    assert not client._owns_executor


@pytest.mark.asyncio
async def test_comfy_mismatch_string_raises():
    """Test that valid config requiring ProcessPoolExecutor raises error if ContextVarExecutor is forced via string."""
    with pytest.raises(ValueError, match="Configuration requires ProcessPoolExecutor"):
        Comfy(configuration={"lowvram": True}, executor="ContextVarExecutor")


@pytest.mark.asyncio
async def test_comfy_mismatch_instance_raises():
    """Test that valid config requiring ProcessPoolExecutor raises error if ContextVarExecutor instance is passed."""
    executor = ContextVarExecutor(max_workers=1)
    with pytest.raises(ValueError, match="Configuration requires ProcessPoolExecutor"):
        Comfy(configuration={"lowvram": True}, executor=executor)


@pytest.mark.asyncio
async def test_comfy_context_manager():
    """Test the async context manager behavior."""
    async with Comfy() as client:
        assert client.is_running
        assert isinstance(client._executor, ContextVarExecutor)
    assert not client.is_running
