# Copyright 2022-2025 ScaleVector
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

"""Helpers for the filesystem resource."""

from typing import Any, Dict, Iterable, List, Optional, Type, Union

import dlt
from dlt.common.configuration import resolve_type
from dlt.common.typing import TDataItem
from dlt.sources import DltResource
from dlt.sources.config import configspec, with_config
from dlt.sources.credentials import (
    CredentialsConfiguration,
    FilesystemConfiguration,
    FileSystemCredentials,
)
from dlt.sources.filesystem import fsspec_filesystem
from fsspec import AbstractFileSystem  # type: ignore


@configspec
class FilesystemConfigurationResource(FilesystemConfiguration):
    credentials: Union[FileSystemCredentials, AbstractFileSystem] = None
    file_glob: Optional[str] = "*"
    files_per_page: int = 100
    extract_content: bool = False

    @resolve_type("credentials")
    def resolve_credentials_type(self) -> Type[CredentialsConfiguration]:
        # use known credentials or empty credentials for unknown protocol
        return Union[
            self.PROTOCOL_CREDENTIALS.get(self.protocol)
            or Optional[CredentialsConfiguration],
            AbstractFileSystem,
        ]  # type: ignore[return-value]


def fsspec_from_resource(filesystem_instance: DltResource) -> AbstractFileSystem:
    """Extract authorized fsspec client from a filesystem resource"""

    @with_config(
        spec=FilesystemConfiguration,
        sections=("sources", filesystem_instance.section, filesystem_instance.name),
    )
    def _get_fsspec(
        bucket_url: str, credentials: Optional[FileSystemCredentials]
    ) -> AbstractFileSystem:
        return fsspec_filesystem(bucket_url, credentials)[0]

    return _get_fsspec(
        filesystem_instance.explicit_args.get("bucket_url", dlt.config.value),
        filesystem_instance.explicit_args.get("credentials", dlt.secrets.value),
    )


def add_columns(columns: List[str], rows: List[List[Any]]) -> List[Dict[str, Any]]:
    """Adds column names to the given rows.

    Args:
        columns (List[str]): The column names.
        rows (List[List[Any]]): The rows.

    Returns:
        List[Dict[str, Any]]: The rows with column names.
    """
    result = []
    for row in rows:
        result.append(dict(zip(columns, row)))

    return result


def fetch_arrow(file_data, chunk_size: int) -> Iterable[TDataItem]:  # type: ignore
    """Fetches data from the given CSV file.

    Args:
        file_data (DuckDBPyRelation): The CSV file data.
        chunk_size (int): The number of rows to read at once.

    Yields:
        Iterable[TDataItem]: Data items, read from the given CSV file.
    """
    batcher = file_data.fetch_arrow_reader(batch_size=chunk_size)
    yield from batcher


def fetch_json(file_data, chunk_size: int) -> List[Dict[str, Any]]:  # type: ignore
    """Fetches data from the given CSV file.

    Args:
        file_data (DuckDBPyRelation): The CSV file data.
        chunk_size (int): The number of rows to read at once.

    Yields:
        Iterable[TDataItem]: Data items, read from the given CSV file.
    """
    while True:
        batch = file_data.fetchmany(chunk_size)
        if not batch:
            break

        yield add_columns(file_data.columns, batch)
