from tongcell.core import ActorId, LocalRuntime

from judy_core.agents.configuration import Configuration
from judy_core.agents.executor import Executor
from judy_core.agents.notification_center import NotificationCenter
from judy_core.agents.output_streamer import OutputStreamer
from judy_core.agents.planner import Planner
from judy_core.agents.state_manager import TaskStateManager
from judy_core.agents.verifier import Verifier
from judy_core.messages import (
    ModelConfig,
    Task,
    TaskConfig,
)


async def create_task_runtime(
    task: Task,
    task_config: TaskConfig,
    model_config: ModelConfig,
) -> tuple[LocalRuntime, OutputStreamer]:
    runtime = LocalRuntime()
    for agent_cls in [Planner, Executor, Verifier, NotificationCenter]:
        await agent_cls.register(runtime)
    config = Configuration(task_config, model_config)
    await runtime.register_actor_factory(Configuration.actor_kind, lambda: config)
    state_manager = TaskStateManager()
    await runtime.register_actor_factory(
        TaskStateManager.actor_kind, lambda: state_manager
    )
    output_streamer = OutputStreamer()
    await runtime.register_actor_factory(
        OutputStreamer.actor_kind, lambda: output_streamer
    )
    await OutputStreamer.register_topic_subscription(runtime)

    await runtime.send(task, ActorId.default(Planner.actor_kind))
    return runtime, output_streamer
