import asyncio

from labels.integrations.docker import get_docker_sbom
from labels.integrations.ecr import ecr_connection, get_token
from labels.model.core import SbomConfig, SourceType
from labels.model.sources import AwsCredentials, AwsRole
from labels.model.syft_sbom import SyftSBOM
from labels.resolvers.directory import Directory
from labels.utils.exceptions import (
    AwsCredentialsNotFoundError,
    AwsExternalIdNotDefinedError,
    AwsRoleNotDefinedError,
    UnexpectedSBOMSourceError,
)


def _resolve_ecr_credentials(role: AwsRole, region: str) -> str:
    return asyncio.run(ecr_connection(role, region=region))


def _resolve_ecr_token(aws_credentials: AwsCredentials, region: str) -> str:
    return asyncio.run(get_token(aws_credentials, region=region))


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


def _handle_docker(*, sbom_config: SbomConfig, is_daemon: bool = False) -> SyftSBOM:
    return get_docker_sbom(
        sbom_config.source,
        username=sbom_config.docker_user,
        password=sbom_config.docker_password,
        os_override=sbom_config.os_override,
        arch_override=sbom_config.arch_override,
        daemon=is_daemon,
    )


def _handle_ecr(sbom_config: SbomConfig) -> SyftSBOM:
    if not sbom_config.aws_role:
        raise AwsRoleNotDefinedError

    if not sbom_config.aws_external_id:
        raise AwsExternalIdNotDefinedError

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

    region = sbom_config.aws_region or "us-east-1"

    aws_creds = _resolve_ecr_credentials(role, region=region)
    return get_docker_sbom(
        sbom_config.source,
        aws_creds=aws_creds,
        os_override=sbom_config.os_override,
        arch_override=sbom_config.arch_override,
    )


def handle_ecr_with_credentials(
    image_uri: str,
    aws_credentials: AwsCredentials,
    *,
    os_override: str | None = None,
    arch_override: str | None = None,
    aws_region: str | None = None,
) -> SyftSBOM:
    region = aws_region or "us-east-1"
    token = _resolve_ecr_token(aws_credentials, region=region)
    return get_docker_sbom(
        image_uri,
        aws_creds=f"AWS:{token}",
        os_override=os_override,
        arch_override=arch_override,
    )


def resolve_sbom_source(sbom_config: SbomConfig) -> Directory | SyftSBOM:
    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 SourceType.ECR_WITH_CREDENTIALS:
            if not sbom_config.aws_credentials:
                raise AwsCredentialsNotFoundError

            return handle_ecr_with_credentials(
                image_uri=sbom_config.source,
                aws_credentials=sbom_config.aws_credentials,
                os_override=sbom_config.os_override,
                arch_override=sbom_config.arch_override,
                aws_region=sbom_config.aws_region,
            )
        case _:
            raise UnexpectedSBOMSourceError(sbom_config.source_type)
