# Copyright 2025 Google LLC
#
# 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.

from __future__ import annotations

import asyncio
import logging
from io import BytesIO
from typing import Any, Dict, List, Optional, Tuple

from google.api_core import exceptions
from google.api_core.retry_async import AsyncRetry
from google.rpc import status_pb2

from google.cloud import _storage_v2
from google.cloud.storage._helpers import generate_random_56_bit_integer
from google.cloud.storage.asyncio._stream_multiplexer import (
    _StreamEnd,
    _StreamError,
    _StreamMultiplexer,
)
from google.cloud.storage.asyncio.async_grpc_client import (
    AsyncGrpcClient,
)
from google.cloud.storage.asyncio.async_read_object_stream import (
    _AsyncReadObjectStream,
)
from google.cloud.storage.asyncio.retry._helpers import _handle_redirect
from google.cloud.storage.asyncio.retry.bidi_stream_retry_manager import (
    _BidiStreamRetryManager,
)
from google.cloud.storage.asyncio.retry.reads_resumption_strategy import (
    _DownloadState,
    _ReadResumptionStrategy,
)
from google.cloud.storage.exceptions import DataCorruption

from ._utils import raise_if_no_fast_crc32c

_MAX_READ_RANGES_PER_BIDI_READ_REQUEST = 100
_BIDI_READ_REDIRECTED_TYPE_URL = (
    "type.googleapis.com/google.storage.v2.BidiReadObjectRedirectedError"
)

logger = logging.getLogger(__name__)


def _is_read_retryable(exc):
    """Predicate to determine if a read operation should be retried."""
    if isinstance(
        exc,
        (
            exceptions.InternalServerError,
            exceptions.ServiceUnavailable,
            exceptions.DeadlineExceeded,
            exceptions.TooManyRequests,
        ),
    ):
        return True

    if not isinstance(exc, exceptions.Aborted) or not exc.errors:
        return False

    try:
        grpc_error = exc.errors[0]
        trailers = grpc_error.trailing_metadata()
        if not trailers:
            return False

        status_details_bin = next(
            (v for k, v in trailers if k == "grpc-status-details-bin"), None
        )

        if not status_details_bin:
            return False

        status_proto = status_pb2.Status()
        status_proto.ParseFromString(status_details_bin)
        return any(
            detail.type_url == _BIDI_READ_REDIRECTED_TYPE_URL
            for detail in status_proto.details
        )
    except Exception as e:
        logger.error(f"Error parsing status_details_bin: {e}")
        return False


class AsyncMultiRangeDownloader:
    """Provides an interface for downloading multiple ranges of a GCS ``Object``
    concurrently.

    Example usage:

    .. code-block:: python

        client = AsyncGrpcClient()
        mrd = await AsyncMultiRangeDownloader.create_mrd(
            client, bucket_name="chandrasiri-rs", object_name="test_open9"
        )
        my_buff1 = open('my_fav_file.txt', 'wb')
        my_buff2 = BytesIO()
        my_buff3 = BytesIO()
        my_buff4 = any_object_which_provides_BytesIO_like_interface()
        await mrd.download_ranges(
            [
                # (start_byte, bytes_to_read, writeable_buffer)
                (0, 100, my_buff1),
                (100, 20, my_buff2),
                (200, 123, my_buff3),
                (300, 789, my_buff4),
            ]
        )

        # verify data in buffers...
        assert my_buff2.getbuffer().nbytes == 20


    """

    @classmethod
    async def create_mrd(
        cls,
        client: AsyncGrpcClient,
        bucket_name: str,
        object_name: str,
        generation: Optional[int] = None,
        read_handle: Optional[_storage_v2.BidiReadHandle] = None,
        retry_policy: Optional[AsyncRetry] = None,
        metadata: Optional[List[Tuple[str, str]]] = None,
        **kwargs,
    ) -> AsyncMultiRangeDownloader:
        """Initializes a MultiRangeDownloader and opens the underlying bidi-gRPC
        object for reading.

        :type client: :class:`~google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient`
        :param client: The asynchronous client to use for making API requests.

        :type bucket_name: str
        :param bucket_name: The name of the bucket containing the object.

        :type object_name: str
        :param object_name: The name of the object to be read.

        :type generation: int
        :param generation: (Optional) If present, selects a specific
                                  revision of this object.

        :type read_handle: _storage_v2.BidiReadHandle
        :param read_handle: (Optional) An existing handle for reading the object.
                            If provided, opening the bidi-gRPC connection will be faster.

        :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry`
        :param retry_policy: (Optional) The retry policy to use for the ``open`` operation.

        :type metadata: List[Tuple[str, str]]
        :param metadata: (Optional) The metadata to be sent with the ``open`` request.

        :rtype: :class:`~google.cloud.storage.asyncio.async_multi_range_downloader.AsyncMultiRangeDownloader`
        :returns: An initialized AsyncMultiRangeDownloader instance for reading.
        """
        mrd = cls(
            client,
            bucket_name,
            object_name,
            generation=generation,
            read_handle=read_handle,
            **kwargs,
        )
        await mrd.open(retry_policy=retry_policy, metadata=metadata)
        return mrd

    def __init__(
        self,
        client: AsyncGrpcClient,
        bucket_name: str,
        object_name: str,
        generation: Optional[int] = None,
        read_handle: Optional[_storage_v2.BidiReadHandle] = None,
        **kwargs,
    ) -> None:
        """Constructor for AsyncMultiRangeDownloader, clients are not adviced to
         use it directly. Instead it's adviced to use the classmethod `create_mrd`.

        :type client: :class:`~google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient`
        :param client: The asynchronous client to use for making API requests.

        :type bucket_name: str
        :param bucket_name: The name of the bucket containing the object.

        :type object_name: str
        :param object_name: The name of the object to be read.

        :type generation: int
        :param generation: (Optional) If present, selects a specific revision of
                                  this object.

        :type read_handle: _storage_v2.BidiReadHandle
        :param read_handle: (Optional) An existing read handle.
        """
        if "generation_number" in kwargs:
            if generation is not None:
                raise TypeError(
                    "Cannot set both 'generation' and 'generation_number'. "
                    "Use 'generation' for new code."
                )
            logger.warning(
                "'generation_number' is deprecated and will be removed in a future "
                "major release. Please use 'generation' instead."
            )
            generation = kwargs.pop("generation_number")

        self.client = client
        self.bucket_name = bucket_name
        self.object_name = object_name
        self.generation = generation
        self.read_handle: Optional[_storage_v2.BidiReadHandle] = read_handle
        self.read_obj_str: Optional[_AsyncReadObjectStream] = None
        self._is_stream_open: bool = False
        self._routing_token: Optional[str] = None
        self._multiplexer: Optional[_StreamMultiplexer] = None
        self.persisted_size: Optional[int] = None  # updated after opening the stream
        self._open_retries: int = 0
        self.is_finalized: bool = False
        self.full_obj_server_crc32c: Optional[int] = None

    async def __aenter__(self):
        """Opens the underlying bidi-gRPC connection to read from the object."""
        await self.open()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Closes the underlying bidi-gRPC connection."""
        if self.is_stream_open:
            await self.close()

    def _on_open_error(self, exc):
        """Extracts routing token and read handle on redirect error during open."""
        logger.warning(f"Error occurred while opening MRD: {exc}")
        routing_token, read_handle = _handle_redirect(exc)
        if routing_token:
            self._routing_token = routing_token
        if read_handle:
            self.read_handle = read_handle

    async def open(
        self,
        retry_policy: Optional[AsyncRetry] = None,
        metadata: Optional[List[Tuple[str, str]]] = None,
    ) -> None:
        """Opens the bidi-gRPC connection to read from the object."""
        if self._is_stream_open:
            raise ValueError("Underlying bidi-gRPC stream is already open")

        if retry_policy is None:

            def on_error_wrapper(exc):
                self._open_retries += 1
                self._on_open_error(exc)

            retry_policy = AsyncRetry(
                predicate=_is_read_retryable, on_error=on_error_wrapper
            )
        else:
            original_on_error = retry_policy._on_error

            def combined_on_error(exc):
                self._open_retries += 1
                self._on_open_error(exc)
                if original_on_error:
                    original_on_error(exc)

            retry_policy = AsyncRetry(
                predicate=_is_read_retryable,
                initial=retry_policy._initial,
                maximum=retry_policy._maximum,
                multiplier=retry_policy._multiplier,
                deadline=retry_policy._deadline,
                on_error=combined_on_error,
            )

        async def _do_open():
            current_metadata = list(metadata) if metadata else []

            # Cleanup stream from previous failed attempt, if any.
            if self.read_obj_str:
                if self.read_obj_str.is_stream_open:
                    try:
                        await self.read_obj_str.close()
                    except exceptions.GoogleAPICallError as e:
                        logger.warning(
                            f"Failed to close existing stream during resumption: {e}"
                        )
                self.read_obj_str = None
                self._is_stream_open = False

            self.read_obj_str = _AsyncReadObjectStream(
                client=self.client.grpc_client,
                bucket_name=self.bucket_name,
                object_name=self.object_name,
                generation_number=self.generation,
                read_handle=self.read_handle,
            )

            if self._routing_token:
                current_metadata.append(
                    ("x-goog-request-params", f"routing_token={self._routing_token}")
                )
                self._routing_token = None

            await self.read_obj_str.open(
                metadata=current_metadata if current_metadata else None
            )

            if self.read_obj_str.generation_number:
                self.generation = self.read_obj_str.generation_number
            if self.read_obj_str.read_handle:
                self.read_handle = self.read_obj_str.read_handle
            if self.read_obj_str.persisted_size is not None:
                self.persisted_size = self.read_obj_str.persisted_size
            self.is_finalized = self.read_obj_str.is_finalized
            self.full_obj_server_crc32c = self.read_obj_str.full_obj_server_crc32c

            self._is_stream_open = True

        await retry_policy(_do_open)()
        self._multiplexer = _StreamMultiplexer(self.read_obj_str)

    def _create_stream_factory(self, state, metadata):
        """Create a factory that opens a new stream with current routing state."""

        async def factory():
            current_handle = state.get("read_handle")
            current_token = state.get("routing_token")

            stream = _AsyncReadObjectStream(
                client=self.client.grpc_client,
                bucket_name=self.bucket_name,
                object_name=self.object_name,
                generation_number=self.generation,
                read_handle=current_handle,
            )

            current_metadata = list(metadata) if metadata else []
            if current_token:
                current_metadata.append(
                    (
                        "x-goog-request-params",
                        f"routing_token={current_token}",
                    )
                )

            await stream.open(metadata=current_metadata if current_metadata else None)

            if stream.generation_number:
                self.generation = stream.generation_number
            if stream.read_handle:
                self.read_handle = stream.read_handle
            self.is_finalized = stream.is_finalized
            self.full_obj_server_crc32c = stream.full_obj_server_crc32c

            self.read_obj_str = stream
            self._is_stream_open = True

            return stream

        return factory

    async def download_ranges(
        self,
        read_ranges: List[Tuple[int, int, BytesIO]],
        lock: asyncio.Lock = None,
        retry_policy: Optional[AsyncRetry] = None,
        metadata: Optional[List[Tuple[str, str]]] = None,
        enable_checksum: bool = True,
    ) -> None:
        """Downloads multiple byte ranges from the object into the buffers
        provided by user with automatic retries.

        :type read_ranges: List[Tuple[int, int, "BytesIO"]]
        :param read_ranges: A list of tuples, where each tuple represents a
            combination of byte_range and writeable buffer in format -
            (`start_byte`, `bytes_to_read`, `writeable_buffer`). Buffer has
            to be provided by the user, and user has to make sure appropriate
            memory is available in the application to avoid out-of-memory crash.

            Special cases:
            if the value of `bytes_to_read` is 0, it'll be interpreted as
            download all contents until the end of the file from `start_byte`.
            Examples:
                * (0, 0, buffer) : downloads 0 to end , i.e. entire object.
                * (100, 0, buffer) : downloads from 100 to end.

        :type lock: asyncio.Lock
        :param lock: (Deprecated) This parameter is deprecated and has no effect.

        :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry`
        :param retry_policy: (Optional) The retry policy to use for the operation.

        :raises ValueError: if the underlying bidi-GRPC stream is not open.
        :raises ValueError: if the length of read_ranges is more than 1000.
        :raises DataCorruption: if a checksum mismatch is detected while reading data.

        """

        if len(read_ranges) > 1000:
            raise ValueError(
                "Invalid input - length of read_ranges cannot be more than 1000"
            )

        if enable_checksum:
            raise_if_no_fast_crc32c()

        if not self._is_stream_open:
            raise ValueError("Underlying bidi-gRPC stream is not open")

        if retry_policy is None:
            retry_policy = AsyncRetry(predicate=_is_read_retryable)

        # Initialize Global State for Retry Strategy
        download_states = {}
        for read_range in read_ranges:
            read_id = generate_random_56_bit_integer()
            # Unpack tuple into self-documenting variable names to improve readability.
            offset, length, user_buffer = read_range

            # Heuristic to detect full object reads:
            # - Implicit full object read: start offset is 0 and length is 0 (read all).
            # - Explicit full object read: start offset is 0 and length matches the exact persisted size.
            is_full_object_read = (offset == 0 and length == 0) or (
                self.persisted_size is not None
                and offset == 0
                and length == self.persisted_size
            )
            download_states[read_id] = _DownloadState(
                initial_offset=offset,
                initial_length=length,
                user_buffer=user_buffer,
                is_full_object_read=is_full_object_read,
            )

        initial_state = {
            "download_states": download_states,
            "read_handle": self.read_handle,
            "routing_token": None,
            "enable_checksum": enable_checksum,
            "full_obj_server_crc32c": self.full_obj_server_crc32c,
        }

        read_ids = set(download_states.keys())
        queue = self._multiplexer.register(read_ids)

        try:
            attempt_count = 0
            last_broken_generation = None

            def send_and_recv_via_multiplexer(
                requests: List[_storage_v2.ReadRange],
                state: Dict[str, Any],
            ):
                async def generator():
                    nonlocal attempt_count, last_broken_generation
                    attempt_count += 1

                    if attempt_count > 1:
                        logger.info(
                            f"Resuming download (attempt {attempt_count}) for {len(requests)} ranges."
                        )

                    # Reopen stream if needed
                    should_reopen = (
                        attempt_count > 1 and last_broken_generation is not None
                    ) or (attempt_count == 1 and metadata is not None)
                    if should_reopen:
                        broken_gen = (
                            last_broken_generation
                            if attempt_count > 1
                            else self._multiplexer.stream_generation
                        )
                        stream_factory = self._create_stream_factory(state, metadata)
                        await self._multiplexer.reopen_stream(
                            broken_gen, stream_factory
                        )

                    stream_generation = self._multiplexer.stream_generation

                    # Send Requests
                    pending_read_ids = {r.read_id for r in requests}
                    for i in range(
                        0, len(requests), _MAX_READ_RANGES_PER_BIDI_READ_REQUEST
                    ):
                        batch = requests[i : i + _MAX_READ_RANGES_PER_BIDI_READ_REQUEST]
                        try:
                            await self._multiplexer.send(
                                _storage_v2.BidiReadObjectRequest(read_ranges=batch)
                            )
                        except Exception:
                            last_broken_generation = stream_generation
                            raise

                    # Receive Responses
                    while pending_read_ids:
                        item = await queue.get()

                        if isinstance(item, _StreamEnd):
                            if pending_read_ids:
                                last_broken_generation = stream_generation
                                raise exceptions.ServiceUnavailable(
                                    "Stream ended with pending read_ids"
                                )
                            break

                        if isinstance(item, _StreamError):
                            if item.generation < stream_generation:
                                continue  # stale error, skip
                            last_broken_generation = item.generation
                            raise item.exception

                        # Track completion
                        if item.object_data_ranges:
                            for data_range in item.object_data_ranges:
                                if data_range.range_end:
                                    pending_read_ids.discard(
                                        data_range.read_range.read_id
                                    )
                        yield item

                return generator()

            strategy = _ReadResumptionStrategy()
            retry_manager = _BidiStreamRetryManager(
                strategy, send_and_recv_via_multiplexer
            )

            try:
                await retry_manager.execute(initial_state, retry_policy)
            except DataCorruption:
                if self.is_stream_open:
                    await self.close()
                raise

            if initial_state.get("read_handle"):
                self.read_handle = initial_state["read_handle"]
        finally:
            if self._multiplexer is not None:
                self._multiplexer.unregister(read_ids)

    async def close(self):
        """
        Closes the underlying bidi-gRPC connection.
        """
        if not self._is_stream_open:
            raise ValueError("Underlying bidi-gRPC stream is not open")

        if self._multiplexer:
            await self._multiplexer.close()
            self._multiplexer = None

        if self.read_obj_str:
            try:
                await self.read_obj_str.close()
            except (asyncio.CancelledError, exceptions.GoogleAPICallError):
                pass
        self.read_obj_str = None
        self._is_stream_open = False

    @property
    def is_stream_open(self) -> bool:
        return self._is_stream_open
