#
# # Copyright © 2026 Peak AI Limited. or its affiliates. All Rights Reserved.
# #
# # Licensed under the Apache License, Version 2.0 (the "License"). You
# # may not use this file except in compliance with the License. A copy of
# # the License is located at:
# #
# # https://github.com/PeakBI/peak-sdk/blob/main/LICENSE
# #
# # or in the "license" file accompanying this file. This file is
# # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# # ANY KIND, either express or implied. See the License for the specific
# # language governing permissions and limitations under the License.
# #
# # This file is part of the peak-sdk.
# # see (https://github.com/PeakBI/peak-sdk)
# #
# # You should have received a copy of the APACHE LICENSE, VERSION 2.0
# # along with this program. If not, see <https://apache.org/licenses/LICENSE-2.0>
#
"""S3 client wrapper class."""

from __future__ import annotations

from typing import Any

import boto3  # type: ignore[import-not-found]
from pydantic import Field  # type: ignore[import-not-found]
from pydantic_settings import BaseSettings  # type: ignore[import-not-found]

from peak.tools.s3.copy import copy
from peak.tools.s3.download import download_file
from peak.tools.s3.exists import object_exists
from peak.tools.s3.upload import upload_file


class S3(BaseSettings):  # type: ignore[misc]
    """S3 client wrapper with common operations."""

    bucket_name: str = Field(..., alias="data_lake")
    root_prefix: str = Field(..., alias="data_lake_root_path")
    s3_session: Any = Field(default_factory=boto3.Session)
    s3_client: Any = Field(default_factory=lambda: boto3.client("s3"))

    class Config:
        """Pydantic configuration."""

        frozen = True

    def object_exists(self, key: str) -> bool:
        """Check if an object exists at the given key."""
        return object_exists(self.s3_client, self.bucket_name, key)

    def copy(
        self,
        source_key: str,
        destination_key: str,
        *,
        exist_ok: bool = True,
    ) -> None:
        """Copy an object from source to destination."""
        copy(self.s3_client, self.bucket_name, source_key, destination_key, exist_ok=exist_ok)

    def download(
        self,
        source_key: str,
        destination: str,
        *,
        exist_ok: bool = False,
    ) -> None:
        """Download an object to a local file."""
        download_file(self.s3_client, self.bucket_name, source_key, destination, exist_ok=exist_ok)

    def upload(
        self,
        destination_key: str,
        file_source: str,
        *,
        exist_ok: bool = False,
    ) -> None:
        """Upload a local file to S3."""
        upload_file(self.s3_client, self.bucket_name, destination_key, file_source, exist_ok=exist_ok)
