from __future__ import annotations

from aicage.registry.custom_agent import load_custom_agents
from aicage.registry.images_metadata.models import ImagesMetadata


def discover_agents(
    images_metadata: ImagesMetadata,
    local_image_repository: str,
) -> ImagesMetadata:
    custom_agents = load_custom_agents(images_metadata, local_image_repository)
    if not custom_agents:
        return images_metadata

    merged_agents = dict(images_metadata.agents)
    merged_agents.update(custom_agents)
    return ImagesMetadata(
        aicage_image=images_metadata.aicage_image,
        aicage_image_base=images_metadata.aicage_image_base,
        bases=images_metadata.bases,
        agents=merged_agents,
    )
