# Copyright (c) 2022 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from time import sleep

from czc.aws_environment import aws_service
from czc.date import get_current_timestamp
from czc.exceptions import CloudWatchQueryException


def query(
    log_group_name: str, query_string: str, query_start_time: int, query_end_time: int = None, limit: int = 100, timeout=30
) -> list[Mapping]:
    end_time = datetime.now(timezone.utc) + timedelta(seconds=timeout)

    query_end_time = query_end_time or int(get_current_timestamp())
    query = aws_service.logs_client.start_query(
        logGroupName=log_group_name, startTime=query_start_time, endTime=query_end_time, queryString=query_string, limit=limit
    )

    query_results = None
    while not query_results or query_results['status'] not in ['Complete', 'Failed', 'Cancelled', 'Timeout']:
        if query_results is not None:
            if datetime.now(timezone.utc) > end_time:
                raise TimeoutError()

            sleep(1)

        query_results = aws_service.logs_client.get_query_results(queryId=query['queryId'])

    if query_results['status'] != 'Complete':
        raise CloudWatchQueryException(query_results['status'])

    return [{entry['field']: entry['value'] for entry in result} for result in query_results['results']]
