import asyncio

from labels.internal.file_resolver.container_image import ContainerImage
from labels.model.core import SbomConfig, SourceType
from labels.sources.directory_source import Directory
from labels.sources.docker import get_docker_image, get_image_context
from labels.sources.ecr import AwsRole, ecr_connection
from labels.utils.exceptions import UnexpectedSBOMSourceError


def _handle_directory(sbom_config: SbomConfig) -> Directory:
    return Directory(
        root=sbom_config.source,
        exclude=sbom_config.exclude,
    )


def _handle_docker(*, sbom_config: SbomConfig, is_daemon: bool = False) -> ContainerImage:
    docker_image = get_docker_image(
        sbom_config.source,
        username=sbom_config.docker_user,
        password=sbom_config.docker_password,
        daemon=is_daemon,
    )

    if not docker_image:
        error_msg = f"No image found for {sbom_config.source}"
        raise ValueError(error_msg)

    context = get_image_context(
        image=docker_image,
        username=sbom_config.docker_user,
        password=sbom_config.docker_password,
        daemon=is_daemon,
    )

    if context is None:
        error_msg = f"No context found for {docker_image}"
        raise ValueError(error_msg)

    return ContainerImage(
        img=docker_image,
        context=context,
        lazy=False,
    )


def _handle_ecr(sbom_config: SbomConfig) -> ContainerImage:
    if not sbom_config.aws_role:
        error_msg = "The AWS role wasn't defined"
        raise ValueError(error_msg)

    role = AwsRole(
        external_id=sbom_config.aws_external_id,
        role=sbom_config.aws_role,
    )

    token, image_metadata = asyncio.run(
        ecr_connection(role, sbom_config.source),
    )

    if not image_metadata:
        error_msg = f"No image found for {sbom_config.source}"
        raise ValueError(error_msg)

    context = get_image_context(
        image=image_metadata,
        aws_creds=f"AWS:{token}",
    )

    if context is None:
        error_msg = f"No context found for {image_metadata}"
        raise ValueError(error_msg)

    return ContainerImage(
        img=image_metadata,
        context=context,
        lazy=False,
    )


def resolve_sbom_source(sbom_config: SbomConfig) -> Directory | ContainerImage:
    match sbom_config.source_type:
        case SourceType.DIRECTORY:
            return _handle_directory(sbom_config)
        case SourceType.DOCKER:
            return _handle_docker(sbom_config=sbom_config)
        case SourceType.DOCKER_DAEMON:
            return _handle_docker(sbom_config=sbom_config, is_daemon=True)
        case SourceType.ECR:
            return _handle_ecr(sbom_config)
        case _:
            raise UnexpectedSBOMSourceError
