import boto3

from src import config
from src.clients.base_client import BaseClient
from src.exceptions import SnowflakeNotFinishedError


class S3Path:
    @staticmethod
    def get_relative_path(data_type: str, run_id: str) -> str:
        """Get S3 path without bucket name.

        Args:
            data_type: Path prefix.
            run_id: Path postfix/part of the current run.

        Returns:
            str: S3 path.

        """
        return f"{config.ENVIRONMENT}/{data_type}/{run_id}/"

    @classmethod
    def get_full_path(cls, data_type: str, run_id: str) -> str:
        """Get full path.

        Args:
            data_type: Path prefix.
            run_id: Path last part of the current run.

        Returns:
            str: Full S3 path.

        """
        return f"s3://{config.S3_BUCKET}/{cls.get_relative_path(data_type, run_id)}"


class S3Client(BaseClient):
    def __init__(self):
        self._resource = boto3.resource("s3", region_name=config.AWS_DEFAULT_REGION)

    def delete_folder_files(self, path: str):
        """Delete all files in folder.

        Args:
            path: S3 path.

        """
        bucket = self._resource.Bucket(config.S3_BUCKET)
        bucket.objects.filter(Prefix=path).delete()

    def count_files(self, path: str) -> int:
        """Count files by prefix.

        Args:
            path: S3 path.

        Returns:
            File count.

        """
        bucket = self._resource.Bucket(config.S3_BUCKET)
        count = 0
        for _ in bucket.objects.filter(Prefix=path).all():
            count += 1
        if not count:
            raise SnowflakeNotFinishedError("S3 files not found")
        return count