import re
import time
from typing import Any, TYPE_CHECKING

import boto3
from botocore.exceptions import ClientError

from aws_testing_utils import config

from .logger import log

if TYPE_CHECKING:
    from mypy_boto3_logs import CloudWatchLogsClient


class CloudwatchLogHandler:
    """Handles interactions with AWS CloudWatch Logs."""

    client: 'CloudWatchLogsClient'
    start_timestamp: int

    def __init__(self) -> None:
        self.client = boto3.client('logs', region_name=config.AWS_REGION)
        now_ns = time.time_ns()
        self.start_timestamp = int(now_ns / 1000000)

    def get_latest_streams(
        self, lambda_name: str, streams_count: int = 1
    ) -> list[dict[str, Any]]:
        """Returns the most recent log streams for a lambda.

        Args:
            lambda_name: Lambda function name in AWS.
            streams_count: Number of streams to return.
        """
        log.info(f'Getting {streams_count} stream(s) for {lambda_name}')
        try:
            stream_response = self.client.describe_log_streams(
                logGroupName=f'/aws/lambda/{lambda_name}',
                orderBy='LastEventTime',
                descending=True,
                limit=streams_count,
            )
            return [dict(s) for s in stream_response['logStreams']]
        except ClientError as exc:
            raise RuntimeError(f'Lambda not found: {lambda_name}') from exc

    def get_stream_logs(self, lambda_name: str, log_stream_name: str) -> dict[str, Any]:
        """Returns log events for a specific lambda stream, filtered to after start_timestamp.

        Args:
            lambda_name: Lambda function name in AWS.
            log_stream_name: CloudWatch log stream name to fetch events from.
        """
        log.info(f'Fetching Logs for Stream Name: {log_stream_name}')
        try:
            logs = self.client.get_log_events(
                logGroupName=f'/aws/lambda/{lambda_name}',
                logStreamName=log_stream_name,
                startTime=self.start_timestamp,
            )
            log.info(f'Lambda logs: {logs}')
            return dict(logs)
        except ClientError as exc:
            raise RuntimeError(
                f'Lambda "{lambda_name}" not found while fetching logs by stream name: "{log_stream_name}"'
            ) from exc

    def get_new_lambda_logs(self, lambda_name: str) -> dict[str, Any]:
        """Returns log events from the latest stream for a lambda, filtered to after start_timestamp."""
        streams = self.get_latest_streams(lambda_name)
        log_stream_name = streams[0]['logStreamName']
        log.info(f'Log Stream Name: {log_stream_name}')

        logs = self.get_stream_logs(lambda_name, log_stream_name)
        return logs

    def assert_lambda_invoked(self, lambda_name: str, seconds_to_wait: int) -> None:
        """Asserts the lambda was invoked after start_timestamp, polling until a log event appears.

        Args:
            lambda_name: Lambda function name in AWS.
            seconds_to_wait: Maximum seconds to poll before raising AssertionError.
        """
        tries = 0
        while tries < seconds_to_wait:
            try:
                logs = self.get_new_lambda_logs(lambda_name)
                if not logs['events']:
                    raise AssertionError
                break
            except AssertionError:
                time.sleep(1)
                tries += 1
        if tries == seconds_to_wait:
            raise AssertionError(
                f'{lambda_name} not invoked after {seconds_to_wait} seconds'
            )

    def assert_lambda_logs_message(
        self,
        lambda_name: str,
        message: str,
        streams_count: int = 10,
        max_tries: int = 20,
    ) -> None:
        """Asserts a message appears in the lambda's logs after start_timestamp.

        Args:
            lambda_name: Lambda function name in AWS.
            message: Regex pattern to search for in log events.
            streams_count: Number of recent streams to search.
            max_tries: Maximum polling attempts (5-second delay between each).
        """
        log.info(f'Start looking for message in cloudwatch logs: "{message}"')
        tries = 0
        is_message = False
        while tries < max_tries:
            log.debug(f'TRY #{tries}')
            streams = self.get_latest_streams(lambda_name, streams_count)
            try:
                for stream in streams:
                    log_stream_name = stream['logStreamName']
                    logs = self.get_stream_logs(lambda_name, log_stream_name)
                    if not logs['events']:
                        log.info('No events in this stream - continue')
                        continue
                    for event in logs['events']:
                        if re.search(message, event['message']):
                            is_message = True
                            log.info('Message found!')
                            break
                    if is_message:
                        break
                assert is_message
                break
            except AssertionError:
                log.info('Nothing found! Sleep and retry')
                time.sleep(5)
                tries += 1
        assert is_message, f'{message} not found in lambda logs'