import pulumi
import pulumi_aws as aws
import conf as config
import json


def create_ecr_repo():
    ecr_repositories = []
    for repo in config.ecr["repos"]:
        if repo["create_ecr_repo"]:
            ecr_repository = aws.ecr.Repository(
                repo,
                name=f"{repo}",
                force_delete=True)

            token = aws.ecr.get_authorization_token_output(registry_id=ecr_repository.registry_id)
            langserve_ecr_life_cycle_policy = aws.ecr.LifecyclePolicy(f"{repo}-ecr-life-cycle-policy",
                repository=ecr_repository.name,
                policy=json.dumps({
                    "rules": [{
                        "rulePriority": 1,
                        "description": "Expire images when they are more than 10 available",
                        "selection": {
                            "tagStatus": "any",
                            "countType": "imageCountMoreThan",
                            "countNumber": 10,
                        },
                        "action": {
                            "type": "expire",
                        },
                    }],
                }))
            
            policy_ecr = aws.iam.get_policy_document(statements=[{
                "sid": "new policy",
                "effect": "Allow",
                "principals": [{
                    "type": "AWS",
                    "identifiers": [config.account_id],
                }],
                "actions": [
                    "ecr:GetDownloadUrlForLayer",
                    "ecr:BatchGetImage",
                    "ecr:BatchCheckLayerAvailability",
                    "ecr:PutImage",
                    "ecr:InitiateLayerUpload",
                    "ecr:UploadLayerPart",
                    "ecr:CompleteLayerUpload",
                    "ecr:DescribeRepositories",
                    "ecr:GetRepositoryPolicy",
                    "ecr:ListImages",
                    "ecr:DeleteRepository",
                    "ecr:BatchDeleteImage",
                    "ecr:SetRepositoryPolicy",
                    "ecr:DeleteRepositoryPolicy",
                ],
            }])
            attach_policy = aws.ecr.RepositoryPolicy(f"{repo}-policy_ecr",
                repository=ecr_repository.name,
                policy=policy_ecr.json)
        else:
            ecr_repository = aws.ecr.get_repository_output(name=repo['name'])
            token = aws.ecr.get_authorization_token_output(registry_id=ecr_repository.registry_id)

        repo['ecr_repo_resource'] = ecr_repository
        repo['ecr_token'] = token
        ecr_repositories.append(repo)
    
    return ecr_repositories

def get_image(ecr_repo_name, image_tag=None, image_digest=None):
    assert (image_tag is not None) != (image_digest is not None), 'User either tag or image_digest, not both, to identify ECR image version.'
    if image_tag:
        return aws.ecr.get_image(repository_name=ecr_repo_name, image_tag=image_tag)
    elif image_digest:
        return aws.ecr.get_image(repository_name=ecr_repo_name, image_digest=image_digest)

def build_and_push(ecr_repositories):
    ecr_repo_images = {}
    for repo in ecr_repositories:
        ecr_repo = repo['ecr_repo_resource']
        container_context = config.get("container-context")
        if container_context is None:
            container_context = "."
        container_file = config.get("container-file")
        if container_file is None:
            container_file = "./Dockerfile"
        
        assert ('tag' in repo.keys()) != ('image_digest' in repo.keys()), 'User must provide either tag or image_digest, but not both, to identify image version'
        if 'tag' in repo.keys():
            ecr_image=aws.ecr.get_image(repository_name=ecr_repo.name, image_tag=repo['tag'])
        elif 'image_digest' in repo.keys():
            ecr_image=aws.ecr.get_image(repository_name=ecr_repo.name, image_digest=repo['image_digest'])
            
        repo['ecr_image'] = ecr_image
        
        ecr_repo_images[repo['name']] = repo

        #ecr_repo_images.append(repo)
    return ecr_repo_images