"""S3 connector."""

import boto3

from src.exceptions import SpatialAssetNotFound


class S3Connector:
    """S3 connector."""

    def __init__(self, environment: str):
        """Initialize AWS session."""
        self.client = boto3.client("s3")
        self.source_bucket = f"{environment}-dolby-atmos-packager-alpha-input"

    def find_spatial_asset_key(self, display_upc: str, isrc: str) -> str:
        """Resolve the full S3 key for a spatial audio file.

        The source bucket is structured as:
            <DISPLAY_UPC>/<DISPLAY_UPC>_<VOLUME>_<TRACK_NUMBER>_<ISRC>.wav

        Prefixes on the UPC folder and matches on the ISRC suffix to
        handle unknown VOLUME and TRACK_NUMBER values.

        Args:
            display_upc (str): spatial display UPC
            isrc (str): spatial ISRC

        Returns:
            (str): S3 key of the spatial asset
        """
        paginator = self.client.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=self.source_bucket, Prefix=f"{display_upc}/{display_upc}_"):
            for obj in page.get("Contents", []):
                key = obj["Key"]
                if key.endswith(f"_{isrc}.wav"):
                    return str(key)

        raise SpatialAssetNotFound(
            f"Spatial ISRC {isrc} was not found on spatial UPC {display_upc} (no matching asset in {self.source_bucket})"
        )