"""async wrapper for asyncio requests"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/client/10_get_data.ipynb.

# %% auto 0
__all__ = ['GetData_Error', 'get_data', 'get_data_stream', 'LooperError', 'looper', 'RouteFunction_ResponseTypeError',
           'route_function']

# %% ../../nbs/client/10_get_data.ipynb 2
from typing import Callable, Optional, Union, Tuple, Any
from functools import wraps
import time

import httpx
import json

from pprint import pprint

import domolibrary.client.DomoAuth as dmda
import domolibrary.client.ResponseGetData as rgd
import domolibrary.client.DomoError as dmde
import domolibrary.client.Logger as dl
import domolibrary.utils.chunk_execution as dmce
import domolibrary.utils.files as dmfi

# %% ../../nbs/client/10_get_data.ipynb 4
class GetData_Error(dmde.DomoError):
    def __init__(self, message, url):
        super().__init__(message=message, domo_instance=url)

# %% ../../nbs/client/10_get_data.ipynb 5
def create_headers(
    auth: dmda.DomoAuth,  # The authentication object containing the Domo API token.
    content_type: dict = None,  # The content type for the request. Defaults to None.
    headers: dict = None,  # Any additional headers for the request. Defaults to None.
) -> dict:  # The headers for the request.
    """
    Creates default headers for interacting with Domo APIs.
    """
    
    if headers is None:
        headers = {}
        
    headers = {
        "Content-Type": content_type or "application/json",
        "Connection": "keep-alive",
        "accept": "application/json, text/plain",
        **headers,
    }
    if auth:
        headers.update(**auth.auth_header)
    return headers

def create_httpx_session(
    session: httpx.AsyncClient = None, 
    is_verify: bool = False
) -> Tuple[httpx.AsyncClient, bool]:
    
    is_close_session = False
    
    if session is None:
        is_close_session = True
        session = httpx.AsyncClient(verify=is_verify)
    return session, is_close_session


# %% ../../nbs/client/10_get_data.ipynb 6
@dmce.run_with_retry()
async def get_data(
    url: str,
    method: str,
    auth: dmda.DomoAuth,
    content_type: Optional[dict] = None,
    headers: Optional[dict] = None,
    body: Union[dict, str, None] = None,
    params: Optional[dict] = None,
    debug_api: bool = False,
    session: httpx.AsyncClient = None,
    return_raw: bool = False,
    is_follow_redirects: bool = False,
    timeout=20,
    parent_class: str = None,  # name of the parent calling class
    num_stacks_to_drop: int = 2,  # number of stacks to drop from the stack trace.  see `domolibrary.client.Logger.TracebackDetails`.  use 2 with class > route structure.  use 1 with route based approach
    debug_traceback: bool = False,
    is_verify: bool = False,
) -> rgd.ResponseGetData:
    """async wrapper for asyncio requests"""

    if debug_api:
        print("🐛 debugging get_data")

    if auth:
        if not auth.token:
            await auth.get_auth_token()

        if not auth.auth_header:
            auth.generate_auth_header()

    headers = create_headers(auth=auth, content_type=content_type, headers=headers)

    session, is_close_session = create_httpx_session(
        session=session, is_verify=is_verify
    )

    traceback_details = dl.get_traceback(
        num_stacks_to_drop=num_stacks_to_drop,
        root_module="<module>",
        parent_class=parent_class,
        debug_traceback=debug_traceback,
    )

    if debug_api:
        pprint(
            {
                "parent_class": parent_class,
                "function_name": traceback_details.function_name,
                "method": method,
                "url": url,
                "headers": headers,
                "body": body,
                "params": params,
            }
        )

    try:
        if isinstance(body, dict) or isinstance(body, list):
            if debug_api:
                print("get_data: sending json")

            if method.lower() == "delete":
                res = httpx.request(
                    method="DELETE",
                    url=url,
                    headers=headers,
                    content=json.dumps(body),
                    params=params,
                    follow_redirects=is_follow_redirects,
                    timeout=timeout,
                )
            else:
                res = await getattr(session, method.lower())(
                    url=url,
                    headers=headers,
                    json=body,
                    params=params,
                    follow_redirects=is_follow_redirects,
                    timeout=timeout,
                )

        elif body:
            if debug_api:
                print("get_data: sending data")

            res = await getattr(session, method.lower())(
                url=url,
                headers=headers,
                data=body,
                params=params,
                follow_redirects=is_follow_redirects,
                timeout=timeout,
            )

        else:
            if debug_api:
                print("get_data: no body")

            res = await getattr(session, method.lower())(
                url=url,
                headers=headers,
                params=params,
                follow_redirects=is_follow_redirects,
                timeout=timeout,
            )

        if debug_api:
            print("get_data_response", res)

        if return_raw:
            return res

        return rgd.ResponseGetData._from_httpx_response(
            res, auth=auth, traceback_details=traceback_details
        )

    finally:
        if is_close_session:
            await session.aclose()

# %% ../../nbs/client/10_get_data.ipynb 11
@dmce.run_with_retry()
async def get_data_stream(
    url: str,
    auth: dmda.DomoAuth,
    method: str = "GET",
    content_type: Optional[dict] = None,
    headers: Optional[dict] = None,
    # body: Union[dict, str, None] = None,
    params: Optional[dict] = None,
    debug_api: bool = False,
    timeout: int = 10,
    parent_class: str = None,  # name of the parent calling class
    num_stacks_to_drop: int = 2,  # number of stacks to drop from the stack trace.  see `domolibrary.client.Logger.TracebackDetails`.  use 2 with class > route structure.  use 1 with route based approach
    debug_traceback: bool = False,
    session: httpx.AsyncClient = None,
    is_verify: bool = False,
    is_follow_redirects: bool = True,
) -> rgd.ResponseGetData:
    """async wrapper for asyncio requests"""

    create_httpx_session(session=session, is_verify=is_verify)
    if debug_api:
        print("🐛 debugging get_data")

    if auth and not auth.token:
        await auth.get_auth_token()

    if headers is None:
        headers = {}

    headers = {
        "Content-Type": content_type or "application/json",
        "Connection": "keep-alive",
        "accept": "application/json, text/plain",
        **headers,
    }

    if auth:
        headers.update(**auth.auth_header)

    traceback_details = dl.get_traceback(
        num_stacks_to_drop=num_stacks_to_drop,
        root_module="<module>",
        parent_class=parent_class,
        debug_traceback=debug_traceback,
    )

    if debug_api:
        pprint(
            {
                "method": method,
                "url": url,
                "headers": headers,
                # "body": body,
                "params": params,
                "traceback_details": traceback_details,
            }
        )


    try:
        async with session or httpx.AsyncClient(verify=False) as client:
            async with client.stream(
                method,
                url=url,
                headers=headers,
                follow_redirects=is_follow_redirects,
            ) as res:
                
                if res.status_code != 200:
                    return rgd.ResponseGetData._from_httpx_response(
                        res = res,
                        auth = auth,
                        parent_class = parent_class,
                        traceback_details = traceback_details
                    )

                content = bytearray()
                async for chunk in res.aiter_bytes():
                    content += chunk
            
                return rgd.ResponseGetData(
                    status = res.status_code,
                    response= content,
                    is_success = True,
                    auth=auth, 
                    traceback_details=traceback_details,
                    parent_class = parent_class
                )

    except httpx.TransportError as e:        
        raise GetData_Error(url=url, message=e) from e


# %% ../../nbs/client/10_get_data.ipynb 15
class LooperError(dmde.DomoError):
    def __init__(self, loop_stage: str, message):
        super().__init__(f"{loop_stage} - {message}")

# %% ../../nbs/client/10_get_data.ipynb 16
async def looper(
    auth: dmda.DomoAuth,
    session: httpx.AsyncClient,
    url,
    offset_params,
    arr_fn: callable,
    loop_until_end: bool = False,  # usually you'll set this to true.  it will override maximum
    method="POST",
    body: dict = None,
    fixed_params: dict = None,
    offset_params_in_body: bool = False,
    body_fn=None,
    limit=1000,
    skip=0,
    maximum=0,
    debug_api: bool = False,
    debug_loop: bool = False,
    debug_num_stacks_to_drop: int = 1,
    parent_class: str = None,
    timeout: bool = 10,
    wait_sleep: int = 0,
    is_verify: bool = False,
    return_raw: bool = False,
) -> rgd.ResponseGetData:
    is_close_session = False

    session, is_close_session = create_httpx_session(session, is_verify=is_verify)

    allRows = []
    isLoop = True

    res = None

    if maximum and maximum <= limit and not loop_until_end:
        limit = maximum

    while isLoop:
        params = fixed_params or {}

        if offset_params_in_body:
            body.update(
                {offset_params.get("offset"): skip, offset_params.get("limit"): limit}
            )

        else:
            params.update(
                {offset_params.get("offset"): skip, offset_params.get("limit"): limit}
            )

        if body_fn:
            try:
                body = body_fn(skip, limit, body)

            except Exception as e:
                await session.aclose()
                raise LooperError(
                    loop_stage="processing body_fn", message=str(e)
                ) from e

        if debug_loop:
            print(f"\n🚀 Retrieving records {skip} through {skip + limit} via {url}")
            # pprint(params)

        res = await get_data(
            auth=auth,
            url=url,
            method=method,
            params=params,
            session=session,
            body=body,
            debug_api=debug_api,
            timeout=timeout,
            parent_class=parent_class,
            num_stacks_to_drop=debug_num_stacks_to_drop,
        )

        if not res.is_success:
            if is_close_session:
                await session.aclose()

            return res
        
        if return_raw:
            return res 

        try:
            newRecords = arr_fn(res)

        except Exception as e:
            await session.aclose()
            raise LooperError(loop_stage="processing arr_fn", message=str(e)) from e

        allRows += newRecords

        if len(newRecords) == 0:
            isLoop = False

        if maximum and len(allRows) >= maximum and not loop_until_end:
            isLoop = False

        if debug_loop:
            print({"all_rows": len(allRows), "new_records": len(newRecords)})
            print(f"skip: {skip}, limit: {limit}")

        if maximum and skip + limit > maximum and not loop_until_end:
            limit = maximum - len(allRows)

        skip += len(newRecords)
        time.sleep(wait_sleep)

    if debug_loop:
        print(
            f"\n🎉 Success - {len(allRows)} records retrieved from {url} in query looper\n"
        )

    if is_close_session:
        await session.aclose()

    return await rgd.ResponseGetData._from_looper(res=res, array=allRows)

# %% ../../nbs/client/10_get_data.ipynb 17
class RouteFunction_ResponseTypeError(TypeError):
    def __init__(self, result):
        super().__init__(
            f"Expected function to return an instance of ResponseGetData got {type(result)} instead.  Refactor function to return ResponseGetData class"
        )


def route_function(func: Callable[..., Any]) -> Callable[..., Any]:
    """
    Decorator for route functions to ensure they receive certain arguments.
    If these arguments are not provided, default values are used.

    Args:
        func (Callable[..., Any]): The function to decorate.

    Returns:
        Callable[..., Any]: The decorated function.

    The decorated function takes the following arguments:
        *args (Any): Positional arguments for the decorated function.
        parent_class (str, optional): The parent class. Defaults to None.
        debug_num_stacks_to_drop (int, optional): The number of stacks to drop for debugging. Defaults to 1.
        debug_api (bool, optional): Whether to debug the API. Defaults to False.
        session (httpx.AsyncClient, optional): The HTTPX client session. Defaults to None.
        **kwargs (Any): Additional keyword arguments for the decorated function.
    """

    @wraps(func)
    async def wrapper(
        *args: Any,
        parent_class: str = None,
        debug_num_stacks_to_drop: int = 1,
        debug_api: bool = False,
        session: httpx.AsyncClient = None,
        **kwargs: Any,
    ) -> Any:
        result = await func(
            *args,
            parent_class=parent_class,
            debug_num_stacks_to_drop=debug_num_stacks_to_drop,
            debug_api=debug_api,
            session=session,
            **kwargs,
        )

        if not isinstance(result, rgd.ResponseGetData):
            raise RouteFunction_ResponseTypeError(result)

        return result

    return wrapper
