from dataclasses import dataclass
from aws.constants import ENVS, CICD_NAMES, ECR_CONFIG
from aws.env_info import get_env_account , get_env_region , get_env_ecr_url
from github.env import exit_on_error_and_write_summary, get_required_env_var
from actions_logging.app_logging import logger
import boto3


@dataclass
class EcrConfig:
    ecr_url: str
    ecr_region: str
    docker_project: str
    account_id: str

    def __init__(self, env_name: str):
        try:
            if env_name not in list(ENVS.keys()) + CICD_NAMES:
                exit_on_error_and_write_summary(f"env_name {env_name} not found in ENVS and CICD_NAMES, not allowed")
            default_config = ECR_CONFIG.get('default')
            self.ecr_url = default_config.get('ecr_url')
            self.ecr_region = default_config.get('ecr_region')
            self.docker_project = default_config.get('docker_project')
            self.account_id = default_config.get('account_id')
            if env_name in ENVS and get_env_ecr_url(env_name):
                self.ecr_url = get_env_ecr_url(env_name)
                self.ecr_region = get_env_region(env_name)
                try:
                    self.account_id = get_env_account(env_name)
                except Exception as account_id_error:
                    exit_on_error_and_write_summary(f"Error setting account_id for env_name {env_name}: {account_id_error}")
                
                self.account_id = str(ENVS[env_name].get('aws_account', self.account_id))
        except Exception as e:
            exit_on_error_and_write_summary(f"unexpected in EcrConfig init - {e}")


def get_repo_manifest(ecr_client, account_id, image_tag, repo_name):
    try:
        response = ecr_client.batch_get_image(
            registryId=account_id,
            repositoryName=repo_name,
            imageIds=[{'imageTag': image_tag}],
            acceptedMediaTypes=['application/vnd.docker.distribution.manifest.v2+json']
        )
        logger.info_green(f"got response for batch_get_image for {repo_name} with image tag {image_tag}")
        if not response['images']:
            exit_on_error_and_write_summary(
                f"no images found in get_repo_manifest for {repo_name} with tag {image_tag}")
        manifest = response["images"][0]["imageManifest"]
        return manifest
    except ecr_client.exceptions.RepositoryNotFoundException as e:
        exit_on_error_and_write_summary(f"repository not found in get_repo_manifest for {repo_name}: {e}")
    except ecr_client.exceptions.ImageNotFoundException as e:
        exit_on_error_and_write_summary(f"image not found in get_repo_manifest for {repo_name}: {e}")
    except Exception as e:
        exit_on_error_and_write_summary(f"error in get_repo_manifest for {repo_name}: {e}")


def update_ecr_tag(current_tag, new_tag, repo_name, env_name="", session_profile=""):
    if not env_name:
        env_name = get_required_env_var("ENV_NAME")
    if session_profile:  # for local usage
        session = boto3.Session(profile_name=session_profile)
    else:
        session = boto3.Session()
    ecr_config = EcrConfig(env_name)
    client = session.client('ecr', region_name=ecr_config.ecr_region)
    manifest = get_repo_manifest(
        ecr_client=client,
        account_id=ecr_config.account_id,
        image_tag=current_tag,
        repo_name=repo_name
    )
    try:
        client.put_image(
            registryId=ecr_config.account_id,
            repositoryName=repo_name,
            imageManifest=manifest,
            imageTag=new_tag,
        )
        logger.info_green(f"success to add tag {new_tag} to {repo_name} with tag {current_tag}")
    except client.exceptions.ImageAlreadyExistsException:
        logger.info_green(f"image tag {new_tag} already exists in {repo_name} for tag {current_tag}")
    except Exception as e:
        exit_on_error_and_write_summary(f"error in update_ecr_tag while trying to add {new_tag} to {repo_name}: {e}")