from __future__ import annotations

from ingestion.common.aws import aws_env
from ingestion.common.aws import list_objects as _s3_list
from ingestion.common.sql import sql_string

from .config import LoadConfig


class S3Storage:
    """StorageAdapter for the shared external loader: lists S3 objects and builds the
    Databend s3:// external stage. Databend reads s3:// natively with access-key credentials
    (the same keys the manifest-driven COPY uses), so no scheme translation is needed."""

    def __init__(self, config: LoadConfig) -> None:
        self._config = config

    def list_objects(self, uri: str) -> list[dict]:
        return _s3_list(uri, env=aws_env(self._config))

    def create_stage_sql(self, stage_name: str, uri: str) -> str:
        connection = [
            f"ACCESS_KEY_ID = {sql_string(self._config.s3_access_key_id)}",
            f"SECRET_ACCESS_KEY = {sql_string(self._config.s3_secret_access_key)}",
            f"REGION = {sql_string(self._config.aws_region)}",
        ]
        # Temporary/STS/SSO credentials carry a session token.
        if self._config.s3_session_token:
            connection.append(f"SESSION_TOKEN = {sql_string(self._config.s3_session_token)}")
        return f"CREATE OR REPLACE STAGE {stage_name} URL = {sql_string(uri)} CONNECTION = ({', '.join(connection)})"

    def credential_env_names(self) -> list[str]:
        return self._config.s3_credential_env_names