from typing import Any, Dict, Optional, Tuple, Union

from lseg_analytics.core.exceptions import LibraryException, ServerError


class AsyncRequestResponse:
    """
    Response object for asynchronous request operations (202 responses).

    This class represents the initial response when an async operation is started,
    containing the necessary information for polling the operation status.

    Attributes:
        status (str): The status of an asynchronous request operation.
            - Received: The request has been successfully received, but is pending processing
            - InProgress: The request is being processed
            - Complete: The request has been completed
        location (str): URL endpoint for polling the operation status
        retry_after (Optional[str]): Indicates how long to wait before polling again
        operation_id (str): Unique identifier for tracking the async request
    """

    def __init__(
        self,
        status: str,
        location: str,
        operation_id: str,
        retry_after: Optional[str] = None,
    ):
        self.status = status
        self.location = location
        self.retry_after = retry_after
        self.operation_id = operation_id


class AsyncPollingResponse:
    """
    Response object for asynchronous polling operations.

    This class represents the response from polling an async operation,

    Attributes:
        status (str): The status of an asynchronous request operation.
               - Received: The request has been successfully received, but is pending processing
               - InProgress: The request is being processed
               - Complete: The request has been completed
        location (Optional[str]): URL endpoint for continued polling (if still in progress). None if completed
        retry_after (Optional[str]): Indicates how long to wait before polling again. None if completed
    """

    def __init__(
        self,
        status: str,
        location: Optional[str] = None,
        retry_after: Optional[str] = None,
    ):
        self.status = status
        self.location = location
        self.retry_after = retry_after


def check_async_polling_response(response: Union[Tuple[Any, Dict[str, str]], Any]) -> AsyncPollingResponse:
    """
    Check async polling response and return structured AsyncPollingResponse.

    Args:
        response: A tuple containing (body, headers) for async polling if not complete, or a direct response object if complete

    Returns:
        AsyncPollingResponse: Structured response with polling information

    Raises:
        LibraryException: If response format is unexpected or invalid
    """

    if isinstance(response, tuple):
        if len(response) != 2:
            raise LibraryException(
                f"Unexpected tuple length for polling response: " f"expected 2 elements, got {len(response)}"
            )
        body, headers = response

        if "location" in headers:  # 200 response (completed)
            return AsyncPollingResponse(
                status=body.status,
                location=headers.get("location"),
                retry_after=None,
            )
        else:  # 202 response (in progress)
            return AsyncPollingResponse(
                status=body.status,
                location=None,
                retry_after=headers.get("retry-after"),
            )


def check_async_request_response(response: Union[Tuple[Any, Dict[str, str]], Any]) -> Union[AsyncRequestResponse, Any]:
    """
    Check async request response and return appropriate response object.

    Args:
        response: Either a tuple (body, headers) for 202 responses (async in progress)
                 or direct response object for 200 responses (completed)

    Returns:
        AsyncRequestResponse: For 202 responses with polling information
        Any: Response data for completed 200 responses

    Raises:
        LibraryException: If response format is unexpected or invalid
    """
    if isinstance(response, tuple):  # 202 response
        if len(response) != 2:
            raise LibraryException(
                f"Unexpected tuple length for async request response: " f"expected 2 elements, got {len(response)}"
            )
        body, headers = response
        return AsyncRequestResponse(
            status=body.status,
            location=headers.get("location"),
            retry_after=headers.get("retry-after"),
            operation_id=headers.get("x-ap-requestid"),
        )

    # This is a 200 response (completed) - return the data directly
    if hasattr(response, "data"):
        return response.data
    else:
        raise ServerError(
            f"Unexpected response type for async request: expected data attribute, got {type(response).__name__}"
        )
