# 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 datetime import datetime, timedelta, timezone
from time import sleep
from typing import List, Mapping

import boto3

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

logs_client = boto3.client('logs')


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 = 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 = 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']]
