from itertools import chain
from aiohttp import ClientResponseError
import asyncio
import aiohttp
import traceback

import demistomock as demisto
import urllib3
from CommonServerPython import *  # noqa # pylint: disable=unused-wildcard-import

# Disable insecure warnings
urllib3.disable_warnings()  # pylint: disable=no-member

""" CONSTANTS """

ALL_SUPPORTED_EVENT_TYPES = ["application", "alert", "page", "audit", "network", "incident"]
MAX_EVENTS_PAGE_SIZE = 10000
MAX_RETRY = 3
NETSKOPE_SEMAPHORE_COUNT = 4
MAX_FAILURE_ENTRIES_TO_HANDLE_PER_TYPE = 10

# Netskope response constants
RATE_LIMIT_REMAINING = "ratelimit-remaining"  # Rate limit remaining
RATE_LIMIT_RESET = "ratelimit-reset"  # Rate limit RESET value is in seconds
VENDOR = "netskope"
PRODUCT = "netskope"

# Event type configuration mapping
# Each event type can have specific endpoint, time parameters, and count field configurations
EVENT_TYPE_CONFIGS: dict[str, dict[str, Any]] = {
    "incident": {
        "endpoint": "/events/datasearch/incident",
        "time_params": {"start_time": "starttime", "end_time": "endtime"},
        "count_field": "event_count:count(_id)",
    },
    # Audit doesn't support the count() aggregation (always returns 0), so skip the count and page
    # directly.
    "audit": {
        "endpoint": "/events/data/{type}",
        "time_params": {"start_time": "insertionstarttime", "end_time": "insertionendtime"},
        "supports_count": False,
    },
}

# Default configuration for all other event types
DEFAULT_EVENT_TYPE_CONFIG: dict[str, Any] = {
    "endpoint": "/events/data/{type}",
    "time_params": {"start_time": "insertionstarttime", "end_time": "insertionendtime"},
    "count_field": "event_count:count(id)",
    "supports_count": True,
}


def get_event_type_config(event_type: str) -> dict:
    """Get configuration for a specific event type.

    Args:
        event_type (str): The type of event

    Returns:
        dict: Configuration dictionary for the event type
    """
    return EVENT_TYPE_CONFIGS.get(event_type, DEFAULT_EVENT_TYPE_CONFIG)


""" CLIENT CLASS """


class Client:
    """
    Client for Netskope RESTful API.

    Args:
        base_url (str): The base URL of Netskope.
        token (str): The token to authenticate against Netskope API.
        validate_certificate (bool): Specifies whether to verify the SSL certificate or not.
        proxy (bool): Specifies if to use XSOAR proxy settings.
    """

    def __init__(self, base_url: str, token: str, proxy: bool, verify: bool, event_types_to_fetch: list[str]):
        self.fetch_status: dict = {event_type: False for event_type in event_types_to_fetch}
        self.event_types_to_fetch: list[str] = event_types_to_fetch
        # Bounds concurrent HTTP requests to Netskope (used inside get_events_data_async).
        self.netskope_semaphore = asyncio.Semaphore(NETSKOPE_SEMAPHORE_COUNT)
        # Bounds the whole page lifecycle (fetch -> send -> free). Separate object from
        # netskope_semaphore to avoid a re-entrant deadlock, so peak memory stays bounded by volume.
        self.page_semaphore = asyncio.Semaphore(NETSKOPE_SEMAPHORE_COUNT)
        self._headers = {"Netskope-Api-Token": f"{token}", "Accept": "application/json"}
        self._base_url = base_url
        self._verify = verify
        self._proxy_url = handle_proxy().get("http") if proxy else None
        self._async_session: aiohttp.ClientSession | None = None

    async def __aenter__(self):
        demisto.debug("Opening the aiohttp session")
        self._async_session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(verify_ssl=self._verify))
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        demisto.debug("Closing aiohttp session")
        if self._async_session is not None:
            await self._async_session.close()

    async def get_events_data_async(self, event_type: str, params: dict) -> dict:
        """Fetch events data asynchronously from Netskope API.

        Args:
            event_type (str): The type of events to fetch (e.g., 'alert', 'network', 'incident')
            params (dict): Query parameters for the API request

        Returns:
            dict: JSON response from the API containing events data

        Raises:
            aiohttp.ClientResponseError: If the HTTP request fails
        """
        # Use the correct endpoint depending on the event type
        config = get_event_type_config(event_type)
        endpoint = config["endpoint"].replace("{type}", event_type)
        url = urljoin(self._base_url, endpoint)
        # Type check for mypy
        if self._async_session is None:
            raise RuntimeError("ClientSession not initialized. Use 'async with' context manager.")
        async with (
            self.netskope_semaphore,
            self._async_session.get(url, params=params, headers=self._headers, proxy=self._proxy_url) as resp,
        ):
            demisto.debug(f"Fetching {event_type} events with params: {params}")
            resp.raise_for_status()
            return await resp.json()

    async def get_events_count(self, event_type: str, params: dict) -> int:
        """Get the count of events for a given type and time range.

        Args:
            event_type (str): The type of events to count
            params (dict): Query parameters for the API request

        Returns:
            int: The count of events available for the given type and time range

        Raises:
            aiohttp.ClientResponseError: If the HTTP request fails
        """
        # Use the correct count field depending on the event type
        config = get_event_type_config(event_type)
        count_field = config["count_field"]

        try:
            res = await self.get_events_data_async(event_type, params | {"fields": count_field})

            # Extract event count from response
            event_count = 0
            if res.get("result") and len(res["result"]) > 0:
                event_count = res["result"][0].get("event_count", 0)

            # Ensure event_count is always a valid integer
            if not isinstance(event_count, int):
                demisto.debug(f"Invalid event_count received: {event_count}, defaulting to 0")
                event_count = 0

            demisto.debug(f"Found {event_count} total {event_type} events for the given time range")
            return event_count

        except Exception as e:
            demisto.error(f"Failed to get event count for {event_type}: {str(e)}")
            raise


""" HELPER FUNCTIONS """


def next_trigger_time(num_of_events, max_fetch, new_last_run):
    """Check whether to add the next trigger key to the next_run dict based on number of fetched events.

    Args:
        num_of_events (int): The number of events fetched.
        max_fetch (int): The maximum fetch limit.
        new_last_run (dict): the next_run to update
    """
    if num_of_events > (max_fetch / 2):
        new_last_run["nextTrigger"] = "0"
    else:
        new_last_run.pop("nextTrigger", None)


def populate_parsing_rule_fields(event: dict, event_type: str):
    """
    Handles the source_log_event and _time fields.
    Sets the source_log_event to the given event type and _time to the time taken from the timestamp field

    Args:
        event (dict): the event to edit
        event_type (str): the event type to set in the source_log_event field
    """
    event["source_log_event"] = event_type
    try:
        event["_time"] = timestamp_to_datestring(event["timestamp"] * 1000, is_utc=True)
    except (TypeError, KeyError):
        # modeling rule will default on ingestion time if _time is missing
        pass


def prepare_events(events: list, event_type: str) -> list:
    """
    Iterates over a list of given events and add/modify special fields like event_id, _time and source_log_event.

    Args:
        events (list): list of events to modify.
        event_type (str): the type of events given in the list.

    Returns:
        list: the list of modified events
    """
    for event in events:
        populate_parsing_rule_fields(event, event_type)
        event_id = event.get("_id")
        event["event_id"] = event_id

    return events


def handle_event_types_to_fetch(event_types_to_fetch) -> list[str]:
    """Handle event_types_to_fetch parameter.
    Transform the event_types_to_fetch parameter into a pythonic list with lowercase values.
    """
    return argToList(
        arg=event_types_to_fetch if event_types_to_fetch else ALL_SUPPORTED_EVENT_TYPES,
        transform=lambda x: x.lower(),
    )


def remove_unsupported_event_types(last_run_dict: dict, event_types_to_fetch: list):
    keys_to_remove = []

    for key in last_run_dict:
        if (key in ALL_SUPPORTED_EVENT_TYPES) and (key not in event_types_to_fetch):
            keys_to_remove.append(key)

    for key in keys_to_remove:
        last_run_dict.pop(key, None)


def get_time_window_params(event_type: str, start_time: str, end_time: str) -> dict:
    """Get time window parameters based on event type configuration.

    Args:
        event_type (str): The type of event
        start_time (str): Start time for the query
        end_time (str): End time for the query

    Returns:
        dict: Time parameters formatted for the specific event type
    """
    config = get_event_type_config(event_type)
    time_params = config["time_params"]
    return {time_params["start_time"]: start_time, time_params["end_time"]: end_time}


def handle_errors(failures):
    failures_res = []
    if failures:
        for failure in failures:
            if isinstance(failure, DemistoException):
                if "Unauthorized" in failure.message:
                    raise failure

                failure_data: dict = failure.res
                demisto.debug(f"error occurred when fetching {failure_data}, {str(failure.exception)}")
                failures_res.append(
                    {
                        "start_time": failure_data.get("insertionstarttime"),
                        "end_time": failure_data.get("insertionendtime"),
                        "offset": failure_data.get("offset", 0),
                        "limit": failure_data.get("limit"),
                    }
                )
            else:
                demisto.error(f"error occurred when fetching, {str(failure)}")

    return failures_res


def handle_prev_fetch_failures(client: Client, last_run: dict, event_type: str, send_to_xsiam: bool, coord_id: str):
    tasks = []
    failure_data = demisto.get(last_run, f"{event_type}.failures", defaultParam=[])
    if failure_data:
        # each failure entry are with the structure {'start_time': ..., 'end_time': ..., 'offset': ..., 'limit':...}
        demisto.debug(
            f"[{coord_id}] there is {len(failure_data)} failure records for {event_type=}, {failure_data=}, handle them"
        )
        for failure_entry in failure_data:
            tasks.append(
                handle_event_type_async(
                    client=client,
                    event_type=event_type,
                    send_to_xsiam=send_to_xsiam,
                    coord_id=coord_id,
                    is_re_fetch_failed_fetch=True,
                    **failure_entry,
                )
            )
    return tasks


async def honor_rate_limiting_async(headers, event_type, params) -> bool:
    """
    Identify the response headers carrying the rate limiting value.
    If the rate limit remaining is 0 then wait for the rate limit reset time before sending the response to the
    client.
    """
    try:
        if RATE_LIMIT_REMAINING in headers:
            remaining = headers.get(RATE_LIMIT_REMAINING)
            demisto.debug(f"Remaining rate limit is: {remaining}")
            if int(remaining) <= 0:
                demisto.debug(f"Rate limiting reached for {event_type=} and {params=}")
                if to_sleep := headers.get(RATE_LIMIT_RESET):
                    demisto.debug(f"Going to async sleep for {to_sleep} seconds to avoid rate limit error")
                    await asyncio.sleep(int(to_sleep))
                else:
                    # if the RESET value does not exist in the header then
                    # sleep for default 1 second as the rate limit remaining is 0
                    demisto.debug("Did not find a rate limit reset value, going to sleep for 1 second to avoid rate limit error")
                    await asyncio.sleep(1)

                return True

    except ValueError as ve:
        logging.error(f"Value error when honoring the rate limiting wait time {headers} {str(ve)}")

    return False


async def handle_event_type_async(
    client: Client,
    event_type: str,
    start_time: str,
    end_time: str,
    offset: int,
    limit: int,
    send_to_xsiam: bool,
    coord_id: str,
    is_re_fetch_failed_fetch: bool = False,
) -> tuple[str, dict]:
    page_size = min(limit, MAX_EVENTS_PAGE_SIZE)
    params = assign_params(limit=page_size, offset=offset, **get_time_window_params(event_type, start_time, end_time))

    demisto.debug(f"[Fetch][{coord_id}] Fetching '{event_type}' events with params: {params}")
    # If this is a retry of a previous failure, log it.
    if is_re_fetch_failed_fetch:
        demisto.debug(f"[Fetch][{coord_id}] Retrying failed fetch for type={event_type}, params={params}")

    success_res, failures = await fetch_and_send_events_async(
        client, event_type, params, limit, send_to_xsiam, is_re_fetch_failed_fetch
    )
    demisto.debug(
        f"[Fetch][{coord_id}] Finished fetching type={event_type} - success pages: {len(success_res)}, failures: {len(failures)}"
    )

    if is_re_fetch_failed_fetch and success_res:
        demisto.debug(f"[Fetch][{coord_id}] Retry succeeded for type={event_type}, params={params}")

    if not success_res and failures and not is_re_fetch_failed_fetch:
        # if there are no success fetch/send, raise an exception and keep the previous next_fetch_start_time
        e: DemistoException = failures[0]
        demisto.error(f"[Fetch][{coord_id}] Failed to fetch events for type={event_type}, params={params}: {str(e)}")
        if hasattr(e, "exception") and hasattr(e.exception, "status"):
            demisto.error(f"[Fetch][{coord_id}] HTTP status: {getattr(e.exception, 'status', None)}")
        if hasattr(e, "res"):
            demisto.error(f"[Fetch][{coord_id}] API response: {getattr(e, 'res', None)}")
        # Try to get status code first, fall back to message checking
        status_code = None
        if hasattr(e, "exception") and hasattr(e.exception, "status"):
            status_code = getattr(e.exception, "status", None)

        # Handle based on status code if available, otherwise use message strings
        error_message = str(e)
        if status_code == 401 or "Unauthorized" in e.message:
            msg = "Unauthorized Error: please validate your credentials and API token permissions."
        elif status_code == 403 or (status_code is None and "forbidden" in error_message.lower()):
            msg = "Forbidden Error: API token lacks required permissions. Please check token permissions in Netskope admin."
        elif "certificate verify failed" in e.message:
            msg = "Connection Error: certificate verification failed, try to use the insecure checkbox."
        elif "Cannot connect to host" in e.message:
            msg = "Connection Error: please validate your Server URL."
        else:
            status_info = f" (HTTP {status_code})" if status_code else ""
            msg = (
                f"Fetching event_type_start_time_end_time_offset: {event_type}_{start_time}_{end_time}_{offset} "
                f"Failed{status_info}, {str(failures[0])}"
            )
        demisto.error(f"[Fetch][{coord_id}] {msg}")
        raise DemistoException(msg, exception=failures[0])
    failures_data = handle_errors(failures)

    # success_res items are ints (per-page counts, fetch path) or event lists (get-events/test path).
    if success_res and isinstance(success_res[0], int):
        # fetch path: pages were already sent & freed, so we only have counts.
        events = []
        events_count = sum(success_res)
    else:
        events = list(chain.from_iterable(success_res))
        events_count = len(events)

    res_dict = {"events": events, "events_count": events_count, "failures": failures_data}
    demisto.debug(f"[Fetch][{coord_id}] Fetched {events_count} {event_type} events ({len(failures_data)} failures)")
    if not is_re_fetch_failed_fetch:
        # if we are retrying a failed fetch (is_re_fetch_failed_fetch=True)
        # no additional info is needed, as we are only trying to fetch the same chunk again.
        if events_count == limit:
            # meaning, there may be another events to fetch for the current time "window"
            # save the start_time and end_time and the next offset

            next_fetch_data = {
                "next_fetch_start_time": start_time,
                "next_fetch_end_time": end_time,
                "next_fetch_offset": offset + events_count,
            }
            demisto.debug(
                f"[Fetch][{coord_id}] reached limit for type={event_type}, storing time window + offset "
                f"for next fetch: {next_fetch_data=}"
            )
            res_dict |= next_fetch_data
        else:
            res_dict |= {"next_fetch_start_time": end_time}

    demisto.debug(f"[Fetch][{coord_id}] Completed type={event_type}, returning {events_count} events")
    return event_type, res_dict


async def fetch_and_send_events_async(
    client: Client, type: str, request_params: dict, limit: int, send_to_xsiam: bool, is_re_fetch_failed_fetch: bool = False
) -> tuple[list, list]:
    """Fetch all pages for a single event type.

    When ``send_to_xsiam`` is True (fetch-events), each page is sent and freed, returning only the
    page count (send-and-flush). When False (get-events), the events are returned for display.

    Returns:
        tuple(success, failures): ``success`` items are ints (counts) when sending, else event lists.
    """

    async def _handle_page(params):
        async def _fetch_page():
            retry_count = 0
            while retry_count < MAX_RETRY:
                try:
                    if retry_count > 0:
                        demisto.debug(f"[Fetch] Rate limit (429) for {type=} {params=}, retry {retry_count=} < {MAX_RETRY=}")

                    offset = params.get("offset")
                    demisto.debug(f"[Fetch] Fetching {type=} page from {offset=}")
                    return await client.get_events_data_async(type, params)
                except ClientResponseError as e:
                    if e.status != 429:  # not rate limit
                        raise e
                    if await honor_rate_limiting_async(e.headers, type, params):
                        retry_count += 1
                    else:
                        raise e
            # Exhausted retries on a 429 - raise so the page is recorded as a failure (and retried
            # next cycle).
            raise DemistoException(f"Rate limit (429) for {type=} not resolved after {MAX_RETRY=} retries")

        async def _send_page_to_xsiam(events):
            # use_streaming_send=True streams+gzips one event at a time (consumes `events`), keeping memory flat.
            demisto.debug(f"[Fetch] Sending {len(events)} {type} events to XSIAM")
            await asyncio.to_thread(
                send_events_to_xsiam,
                events=events,
                vendor=VENDOR,
                product=PRODUCT,
                chunk_size=XSIAM_EVENT_CHUNK_SIZE_LIMIT,
                use_streaming_send=True,
            )

        # Bound the whole page lifecycle (fetch -> send -> free) so at most NETSKOPE_SEMAPHORE_COUNT
        # pages are in memory at once across all types, keeping peak memory independent of volume.
        async with client.page_semaphore:
            try:
                res = await _fetch_page()
                events = res.get("result", [])
                events = prepare_events(events, type)
                page_count = len(events)
                demisto.debug(f"[Fetch] Fetched {page_count} {type} events for page (offset={params.get('offset')})")
                if send_to_xsiam:
                    # stream-send then DROP the page: return only the count so nothing accumulates upstream
                    await _send_page_to_xsiam(events)
                    demisto.debug(f"[Fetch] Sent and freed {page_count} {type} events (send-and-flush)")
                    return page_count
            except Exception as e:
                demisto.error(f"[Fetch] Error handling {type} page (offset={params.get('offset')}): {str(e)}")
                raise DemistoException(message=str(e), exception=e, res=params)
            # get-events path: caller needs the actual events
            return events

    def _page_result_len(page_result: int | list[dict]) -> int:
        # _handle_page returns an int (fetch-events: streamed & freed) or the events list (get-events).
        return page_result if isinstance(page_result, int) else len(page_result)

    async def _handle_all_pages_sequential():
        """Paginate without a pre-flight count (for datasets like `audit` that don't support count()).

        Pages one at a time until a short page (last page) or `limit` events. Returns success items
        and BaseException items (like return_exceptions=True) so failures are recorded, not swallowed.
        """
        init_offset = int(request_params.pop("offset", 0))
        request_limit = int(request_params.get("limit", MAX_EVENTS_PAGE_SIZE))
        max_offset = init_offset + int(limit)
        demisto.debug(
            f"[Fetch] type={type}: sequential paging (count unsupported) from {init_offset=} "
            f"up to {limit} events in pages of {request_limit}"
        )

        results: list = []
        offset = init_offset
        while offset < max_offset:
            # Cap total at `limit` so the caller's `events_count == limit` check checkpoints next_fetch_offset.
            page_size = min(request_limit, max_offset - offset)
            try:
                page_result = await _handle_page(request_params | {"offset": offset, "limit": page_size})
            except Exception as e:
                # Record the failure (don't abort the window) so the caller checkpoints and retries this offset.
                demisto.error(f"[Fetch] type={type}: sequential page failed at {offset=}: {str(e)}")
                demisto.debug(traceback.format_exc())
                results.append(e)
                break
            results.append(page_result)
            if _page_result_len(page_result) < page_size:
                break  # short page => last page for this window
            offset += page_size
        return results

    async def _handle_all_pages():
        try:
            # Datasets that don't support the count aggregation must page directly (no pre-flight count).
            if not is_re_fetch_failed_fetch and not get_event_type_config(type).get("supports_count", True):
                return await _handle_all_pages_sequential()

            # the `offset` should not be in the get_events_count request
            init_offset = int(request_params.pop("offset", 0))

            if is_re_fetch_failed_fetch:
                # in case of re-fetch failures we don't use pagination, just the fetch the failed chunk
                total_events = limit
                max_offset = init_offset + int(limit)
            else:
                total_events = await client.get_events_count(type, request_params)
                max_offset = min(total_events, init_offset + int(limit))

            request_limit = request_params.get("limit", MAX_EVENTS_PAGE_SIZE)
            demisto.debug(
                f"[Fetch] type={type}: fetching up to {min(total_events, limit)} events from {init_offset=} "
                f"in pages of {request_limit}"
            )
            tasks = [
                _handle_page(request_params | {"offset": offset}) for offset in range(init_offset, max_offset, request_limit)
            ]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            return results
        except Exception as e:
            raise DemistoException(message=str(e), exception=e, res=request_params)

    try:
        results: list[list[dict] | BaseException] = await _handle_all_pages()
        success_tasks = list(filter(lambda res: not isinstance(res, BaseException), results))
        failures = list(filter(lambda res: isinstance(res, BaseException), results))
        return success_tasks, failures
    except Exception as e:
        return [], [e]


""" COMMANDS FUNCTIONS """


async def handle_fetch_and_send_all_events(
    client: Client, last_run: dict, limit: int = MAX_EVENTS_PAGE_SIZE, send_to_xsiam=False
) -> tuple[list[dict], int, dict]:
    """
    Iterates over all supported event types and call the handle event fetch logic and send the events to XSIAM.

    Endpoint: /api/v2/events/data/
    Docs: https://www.postman.com/netskope-tech-alliances/netskope-rest-api/request/zknja6y/get-network-events-generated-by-netskope

    Example HTTP request:
    <baseUrl>/api/v2/events/data/network?offset=0&insertionstarttime=1707466628&insertionendtime=1739089028&

    Args:
        client (Client): The Netskope client.
        last_run (dict): The execution last run dict where the relevant operations are stored.
        limit (int): The limit which after we stop pulling.
        send_to_xsiam(bool): Whether to send the fetched events to XSIAM or not.

    Returns:
        list: The accumulated events (empty on the fetch path, where events are sent and freed).
        int: The total number of events fetched.
        dict: The updated last_run object.
    """
    start = time.time()
    # needed as we use concurrent async tasks
    support_multithreading()

    remove_unsupported_event_types(last_run, client.event_types_to_fetch)

    all_events = []
    total_events_count = 0
    epoch_current_time = str(int(arg_to_datetime("now").timestamp()))  # type: ignore[union-attr]
    epoch_last_day = str(int(arg_to_datetime("1 day").timestamp()))  # type: ignore[union-attr]
    page_size = min(limit, MAX_EVENTS_PAGE_SIZE)

    # Create main coordination ID for async logging traceability
    coord_id = f"coord_{int(time.time() * 1000) % 10000}"
    demisto.info(
        f"[Fetch][{coord_id}] Starting fetch cycle for types={client.event_types_to_fetch} "
        f"({page_size=}, {limit=}, send_to_xsiam={send_to_xsiam})"
    )

    prev_fetch_failure_tasks = []
    new_tasks = []
    for event_type in client.event_types_to_fetch:
        # for each event type, we run 2 separated async fetch
        # 1. to fetch the previous failed fetches (which stored in the last run) collected in prev_fetch_failure_tasks
        # 2. fetch the new events (regular fetch) collected in the new_tasks list

        demisto.debug(f"[Fetch][{coord_id}] Processing event type: {event_type}")
        # get failures from previous iteration
        prev_fetch_failure_tasks.extend(handle_prev_fetch_failures(client, last_run, event_type, send_to_xsiam, coord_id))
        last_run_current_type = last_run.get(event_type, {})
        start_time = last_run_current_type.get("next_fetch_start_time", epoch_last_day)
        end_time = last_run_current_type.get("next_fetch_end_time", epoch_current_time)
        offset = int(last_run_current_type.get("next_fetch_offset", 0))
        demisto.debug(
            f"[Fetch][{coord_id}] Scheduling task for {event_type}: start={start_time}, end={end_time}, offset={offset}"
        )
        new_tasks.append(
            handle_event_type_async(client, event_type, start_time, end_time, offset, limit, send_to_xsiam, coord_id)
        )

    demisto.debug(f"[Fetch][{coord_id}] Running {len(prev_fetch_failure_tasks)} retry tasks + {len(new_tasks)} new tasks")
    results = await asyncio.gather(*prev_fetch_failure_tasks, *new_tasks, return_exceptions=True)
    success_tasks = list(filter(lambda res: not isinstance(res, BaseException), results))
    failures_tasks = list(filter(lambda res: isinstance(res, BaseException), results))
    demisto.info(f"[Fetch][{coord_id}] Type tasks done - success: {len(success_tasks)}, failed: {len(failures_tasks)}")

    if failures_tasks and not success_tasks:
        # meaning, all the tasks was failed
        demisto.error(f"[Fetch][{coord_id}] All event-type tasks failed, raising the first exception")
        raise DemistoException(failures_tasks[0])
    new_last_run: dict = {}
    for task_result in success_tasks:
        # Type check for mypy
        if isinstance(task_result, tuple):
            event_type, event_type_res = task_result
            # event_type_res is in structure of:
            # {'events':[...], 'events_count': N, 'failures':[...], additional data like next_run_start_time, next_run_offset}
            total_events_count += int(event_type_res.pop("events_count", 0))
            all_events.extend(event_type_res.pop("events", []))
            existing_failures = demisto.get(new_last_run, f"{event_type}.failures", defaultParam=[])
            existing_failures.extend(event_type_res.pop("failures", []))

            # in the init, set to the old last_run data
            new_last_run.setdefault(event_type, last_run.get(event_type, {}))
            if event_type_res:
                # in case of new data - override the old data
                new_last_run[event_type] = event_type_res
            if len(existing_failures) > MAX_FAILURE_ENTRIES_TO_HANDLE_PER_TYPE:
                demisto.debug(
                    f"[Fetch][{coord_id}] Truncating failures for {event_type}: "
                    f"{len(existing_failures)} > {MAX_FAILURE_ENTRIES_TO_HANDLE_PER_TYPE}, keeping the first ones."
                )
            new_last_run[event_type]["failures"] = existing_failures[:MAX_FAILURE_ENTRIES_TO_HANDLE_PER_TYPE]

    demisto.info(f"[Fetch][{coord_id}] Cycle done: handled {total_events_count} total events in {time.time() - start:.2f}s")

    return all_events, total_events_count, new_last_run


async def get_events_command_async(
    client: Client, args: dict[str, Any], last_run: dict, should_push_events: bool = False
) -> CommandResults:
    """Manual netskope-get-events command: fetch a small batch, optionally push it, and display it.

    Optionally accepts a manual time window via the `start_time`/`end_time` args. When provided,
    we synthesize a `last_run` that pins the window for every fetched type, so the existing fetch
    pipeline (including the per-type API param mapping and the audit no-count path) is reused as-is.
    When the args are omitted, the instance's real `last_run` is used - i.e. the default behavior
    is unchanged.
    """
    limit = arg_to_number(args.get("limit")) or 10

    # Optional manual time window: triggered by start_time (end_time defaults to "now"). Without
    # start_time the instance's real last_run is used (default behavior). For audit this filters by
    # insertion time (insertionstarttime/endtime). `arg_name` makes arg_to_datetime raise a clear
    # error on unparseable input, so the parsed values are always valid datetimes.
    if args.get("start_time"):
        start_arg = arg_to_datetime(args.get("start_time"), arg_name="start_time")
        end_arg = arg_to_datetime(args.get("end_time") or "now", arg_name="end_time")
        start_ts = int(start_arg.timestamp())  # type: ignore[union-attr]
        end_ts = int(end_arg.timestamp())  # type: ignore[union-attr]
        if end_ts <= start_ts:
            return_error(f"'end_time' ({end_ts}) must be after 'start_time' ({start_ts}).")
        start_epoch = str(start_ts)
        end_epoch = str(end_ts)
        last_run = {
            fetch_type: {"next_fetch_start_time": start_epoch, "next_fetch_end_time": end_epoch, "failures": []}
            for fetch_type in client.event_types_to_fetch
        }
        demisto.debug(
            f"[Get-Events] Using manual time window {start_epoch} -> {end_epoch} for types={client.event_types_to_fetch}"
        )

    demisto.debug(f"[Get-Events] Running netskope-get-events with {limit=}, {should_push_events=}")

    # Two distinct flows use send_to_xsiam differently:
    #   - send_to_xsiam (the fetch flow): stream-and-flush each page, return only counts (memory-bounded).
    #   - should_push_events (this get-events flow): we need the events to DISPLAY them, so fetch with
    #     send_to_xsiam=False to get them back, and push separately below only if asked.
    events, _, _ = await handle_fetch_and_send_all_events(client=client, last_run=last_run, limit=limit, send_to_xsiam=False)
    demisto.debug(f"[Get-Events] Fetched {len(events)} events for display")

    # Low-volume command, so push with the plain (non-streaming) send, which keeps `events` intact for display.
    if should_push_events:
        demisto.debug(f"[Get-Events] Pushing {len(events)} events to XSIAM (non-streaming)")
        send_events_to_xsiam(events, vendor=VENDOR, product=PRODUCT)

    for event in events:
        event["timestamp"] = timestamp_to_datestring(event["timestamp"] * 1000)

    readable_output = tableToMarkdown(
        "Events List:",
        events,
        removeNull=True,
        headers=["_id", "timestamp", "type", "access_method", "app", "traffic_type"],
        headerTransform=string_to_table_header,
    )

    results = CommandResults(
        outputs_prefix="Netskope.Event",
        outputs_key_field="_id",
        outputs=events,
        readable_output=readable_output,
        raw_response=events,
    )

    return results


async def test_module(client: Client, last_run: dict) -> str:
    await get_events_command_async(client=client, args={"limit": 1}, last_run=last_run, should_push_events=False)
    return "ok"


""" MAIN FUNCTION """


async def main() -> None:  # pragma: no cover
    try:
        params = demisto.params()

        url = params.get("url")
        base_url = urljoin(url, "/api/v2/")
        token = params.get("credentials", {}).get("password")
        verify_certificate = not params.get("insecure", False)
        proxy = params.get("proxy", False)
        max_fetch: int = arg_to_number(params.get("max_fetch")) or 10000

        command_name = demisto.command()
        demisto.debug(f"Command being called is {command_name}")

        event_types_to_fetch = handle_event_types_to_fetch(params.get("event_types_to_fetch"))
        demisto.debug(f"Event types that will be fetched in this instance: {event_types_to_fetch}")

        async with Client(base_url, token, proxy, verify_certificate, event_types_to_fetch) as client:
            last_run = demisto.getLastRun()
            demisto.debug(f"Running with the following last_run - {last_run}")

            new_last_run: dict = {}
            if command_name == "test-module":
                # This is the call made when pressing the integration Test button.
                result = await test_module(client, last_run)  # type: ignore[arg-type]
                return_results(result)

            elif command_name == "netskope-get-events":
                args = demisto.args()
                should_push_events = argToBoolean(args.get("should_push_events", "true"))
                demisto.debug(f"Running netskope-get-events with should_push_events={should_push_events}")
                results = await get_events_command_async(client, args, last_run, should_push_events=should_push_events)
                return_results(results)

            elif command_name == "fetch-events":
                demisto.debug(f"Starting fetch with last run {last_run}")
                # send-and-flush: events are sent per page and freed; we only get the total count back
                _all_events, total_events_count, new_last_run = await handle_fetch_and_send_all_events(
                    client=client, last_run=last_run, limit=max_fetch, send_to_xsiam=True
                )
                demisto.debug(f"Fetched {total_events_count} total events.")
                next_trigger_time(total_events_count, max_fetch, new_last_run)
                demisto.debug(f"Setting the last_run to: {new_last_run}")
                demisto.setLastRun(new_last_run)

    except Exception as e:
        # Log the specific exception type and full traceback for better debugging
        error_traceback = traceback.format_exc()
        demisto.error(f"{type(e).__name__} in {command_name}: {str(e)}\nTraceback:\n{error_traceback}")
        last_run = new_last_run if new_last_run else demisto.getLastRun()
        last_run.pop("nextTrigger", None)
        demisto.setLastRun(last_run)
        demisto.debug(f"last run after removing nextTrigger {last_run}")
        return_error(f"Failed to execute {command_name} command.\nError: {type(e).__name__}: {str(e)}")


""" ENTRY POINT """

if __name__ in ("__main__", "__builtin__", "builtins"):
    asyncio.run(main())