import base64
import gzip
import os
import random
import string
import sys
from contextlib import contextmanager
from datetime import datetime
from threading import Event, Thread
from typing import (
    IO,
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Generator,
    Iterator,
    List,
    Optional,
    Sequence,
    Tuple,
    cast,
)

import boto3
import dagster._check as check
from botocore.exceptions import ClientError
from dagster import DagsterInvariantViolationError
from dagster._annotations import experimental
from dagster._core.pipes.client import PipesLaunchedData, PipesMessageReader, PipesParams
from dagster._core.pipes.context import PipesMessageHandler
from dagster._core.pipes.utils import (
    PipesBlobStoreMessageReader,
    PipesChunkedLogReader,
    PipesLogReader,
    PipesThreadedMessageReader,
    extract_message_or_forward_to_stdout,
    forward_only_logs_to_file,
)
from dagster._utils.backoff import backoff
from dagster_pipes import PipesBlobStoreMessageWriter, PipesDefaultMessageWriter

if TYPE_CHECKING:
    from mypy_boto3_logs import CloudWatchLogsClient
    from mypy_boto3_logs.type_defs import OutputLogEventTypeDef
    from mypy_boto3_s3 import S3Client


def _can_read_from_s3(client: "S3Client", bucket: Optional[str], key: Optional[str]):
    if not bucket or not key:
        return False
    else:
        try:
            client.head_object(Bucket=bucket, Key=key)
            return True
        except ClientError:
            return False


def default_log_decode_fn(contents: bytes) -> str:
    return contents.decode("utf-8")


def gzip_log_decode_fn(contents: bytes) -> str:
    return gzip.decompress(contents).decode("utf-8")


class PipesS3LogReader(PipesChunkedLogReader):
    def __init__(
        self,
        *,
        bucket: str,
        key: str,
        client: Optional["S3Client"] = None,
        interval: float = 10,
        target_stream: Optional[IO[str]] = None,
        # TODO: maybe move this parameter to a different scope
        decode_fn: Optional[Callable[[bytes], str]] = None,
        debug_info: Optional[str] = None,
    ):
        self.bucket = bucket
        self.key = key
        self.client: "S3Client" = client or boto3.client("s3")
        self.decode_fn = decode_fn or default_log_decode_fn

        self.log_position = 0

        super().__init__(
            interval=interval, target_stream=target_stream or sys.stdout, debug_info=debug_info
        )

    @property
    def name(self) -> str:
        return f"PipesS3LogReader(s3://{os.path.join(self.bucket, self.key)})"

    def target_is_readable(self, params: PipesParams) -> bool:
        return _can_read_from_s3(
            client=self.client,
            bucket=self.bucket,
            key=self.key,
        )

    def download_log_chunk(self, params: PipesParams) -> Optional[str]:
        text = self.decode_fn(
            self.client.get_object(Bucket=self.bucket, Key=self.key)["Body"].read()
        )
        current_position = self.log_position
        self.log_position += len(text)

        return text[current_position:]


class PipesS3MessageReader(PipesBlobStoreMessageReader):
    """Message reader that reads messages by periodically reading message chunks from a specified S3
    bucket.

    If `log_readers` is passed, this reader will also start the passed readers
    when the first message is received from the external process.

    Args:
        interval (float): interval in seconds between attempts to download a chunk
        bucket (str): The S3 bucket to read from.
        client (WorkspaceClient): A boto3 client.
        log_readers (Optional[Sequence[PipesLogReader]]): A set of log readers for logs on S3.
        include_stdio_in_messages (bool): Whether to send stdout/stderr to Dagster via Pipes messages. Defaults to False.
    """

    def __init__(
        self,
        *,
        interval: float = 10,
        bucket: str,
        client: boto3.client,  # pyright: ignore (reportGeneralTypeIssues)
        log_readers: Optional[Sequence[PipesLogReader]] = None,
        include_stdio_in_messages: bool = False,
    ):
        super().__init__(
            interval=interval,
            log_readers=log_readers,
        )
        self.bucket = check.str_param(bucket, "bucket")
        self.include_stdio_in_messages = check.bool_param(
            include_stdio_in_messages, "include_stdio_in_messages"
        )
        self.client = client

    @contextmanager
    def get_params(self) -> Iterator[PipesParams]:
        key_prefix = "".join(random.choices(string.ascii_letters, k=30))
        yield {
            "bucket": self.bucket,
            "key_prefix": key_prefix,
            PipesBlobStoreMessageWriter.INCLUDE_STDIO_IN_MESSAGES_KEY: self.include_stdio_in_messages,
        }

    def messages_are_readable(self, params: PipesParams) -> bool:
        key_prefix = params.get("key_prefix")
        if key_prefix is not None:
            try:
                self.client.head_object(Bucket=self.bucket, Key=f"{key_prefix}/1.json")
                return True
            except ClientError:
                return False
        else:
            return False

    def download_messages_chunk(self, index: int, params: PipesParams) -> Optional[str]:
        key = f"{params['key_prefix']}/{index}.json"
        try:
            obj = self.client.get_object(Bucket=self.bucket, Key=key)
            return obj["Body"].read().decode("utf-8")
        except ClientError:
            return None

    def no_messages_debug_text(self) -> str:
        return (
            f"Attempted to read messages from S3 bucket {self.bucket}. Expected"
            " PipesS3MessageWriter to be explicitly passed to open_dagster_pipes in the external"
            " process."
        )


class PipesLambdaLogsMessageReader(PipesMessageReader):
    """Message reader that consumes buffered pipes messages that were flushed on exit from the
    final 4k of logs that are returned from issuing a sync lambda invocation. This means messages
    emitted during the computation will only be processed once the lambda completes.

    Limitations: If the volume of pipes messages exceeds 4k, messages will be lost and it is
    recommended to switch to PipesS3MessageWriter & PipesS3MessageReader.
    """

    @contextmanager
    def read_messages(
        self,
        handler: PipesMessageHandler,
    ) -> Iterator[PipesParams]:
        self._handler = handler
        try:
            # use buffered stdio to shift the pipes messages to the tail of logs
            yield {PipesDefaultMessageWriter.BUFFERED_STDIO_KEY: PipesDefaultMessageWriter.STDERR}
        finally:
            self._handler = None

    def consume_lambda_logs(self, response) -> None:
        handler = check.not_none(
            self._handler, "Can only consume logs within context manager scope."
        )

        log_result = base64.b64decode(response["LogResult"]).decode("utf-8")

        for log_line in log_result.splitlines():
            extract_message_or_forward_to_stdout(handler, log_line)

    def no_messages_debug_text(self) -> str:
        return (
            "Attempted to read messages by extracting them from the tail of lambda logs directly."
        )


# Number of retries to attempt getting cloudwatch logs when faced with a throttling exception.
DEFAULT_CLOUDWATCH_LOGS_MAX_RETRIES = 10


# Custom backoff delay_generator for get_log_events which adds some jitter
def get_log_events_delay_generator() -> Iterator[float]:
    i = 0.5
    while True:
        yield i
        i *= 2
        i += random.uniform(0, 1)


def get_log_events(
    client: "CloudWatchLogsClient",
    max_retries: Optional[int] = DEFAULT_CLOUDWATCH_LOGS_MAX_RETRIES,
    **log_params,
):
    max_retries = max_retries or DEFAULT_CLOUDWATCH_LOGS_MAX_RETRIES

    return backoff(
        fn=client.get_log_events,
        kwargs=log_params,
        retry_on=(client.exceptions.ThrottlingException,),
        max_retries=max_retries,
        delay_generator=get_log_events_delay_generator(),
    )


def tail_cloudwatch_events(
    client: "CloudWatchLogsClient",
    log_group: str,
    log_stream: str,
    start_time: Optional[int] = None,
    max_retries: Optional[int] = DEFAULT_CLOUDWATCH_LOGS_MAX_RETRIES,
) -> Generator[List["OutputLogEventTypeDef"], None, None]:
    """Yields events from a CloudWatch log stream."""
    params: Dict[str, Any] = {
        "logGroupName": log_group,
        "logStreamName": log_stream,
    }

    if start_time is not None:
        params["startTime"] = start_time

    response = get_log_events(client=client, max_retries=max_retries, **params)

    while True:
        events = response.get("events")

        if events:
            yield events

        params["nextToken"] = response["nextForwardToken"]

        response = get_log_events(client=client, max_retries=max_retries, **params)


@experimental
class PipesCloudWatchLogReader(PipesLogReader):
    def __init__(
        self,
        client=None,
        log_group: Optional[str] = None,
        log_stream: Optional[str] = None,
        target_stream: Optional[IO[str]] = None,
        start_time: Optional[int] = None,
        debug_info: Optional[str] = None,
        max_retries: Optional[int] = DEFAULT_CLOUDWATCH_LOGS_MAX_RETRIES,
    ):
        self.client = client or boto3.client("logs")
        self.log_group = log_group
        self.log_stream = log_stream
        self.target_stream = target_stream or sys.stdout
        self.thread = None
        self.start_time = start_time
        self._debug_info = debug_info
        self.max_retries = max_retries

    @property
    def debug_info(self) -> Optional[str]:
        return self._debug_info

    def target_is_readable(self, params: PipesParams) -> bool:
        log_group = params.get("log_group") or self.log_group
        log_stream = params.get("log_stream") or self.log_stream

        if log_group is not None and log_stream is not None:
            # check if the stream actually exists
            try:
                resp = self.client.describe_log_streams(
                    logGroupName=log_group,
                    logStreamNamePrefix=log_stream,
                )

                if resp.get("logStreams", []):
                    return True
                else:
                    return False
            except self.client.exceptions.ResourceNotFoundException:
                return False
        else:
            return False

    def start(self, params: PipesParams, is_session_closed: Event) -> None:
        if not self.target_is_readable(params):
            raise DagsterInvariantViolationError(
                "log_group and log_stream must be set either in the constructor or in Pipes params."
            )

        self.thread = Thread(
            target=self._start, kwargs={"params": params, "is_session_closed": is_session_closed}
        )
        self.thread.start()

    def _start(self, params: PipesParams, is_session_closed: Event) -> None:
        log_group = cast(str, params.get("log_group") or self.log_group)
        log_stream = cast(str, params.get("log_stream") or self.log_stream)
        start_time = cast(int, self.start_time or params.get("start_time"))

        for events in tail_cloudwatch_events(
            self.client, log_group, log_stream, start_time=start_time, max_retries=self.max_retries
        ):
            for event in events:
                for line in event.get("message", "").splitlines():
                    if line:
                        forward_only_logs_to_file(line, self.target_stream)

            if is_session_closed.is_set():
                return

    def stop(self) -> None:
        pass

    def is_running(self) -> bool:
        return self.thread is not None and self.thread.is_alive()


@experimental
class PipesCloudWatchMessageReader(PipesThreadedMessageReader):
    """Message reader that consumes AWS CloudWatch logs to read pipes messages."""

    def __init__(
        self,
        client=None,
        log_group: Optional[str] = None,
        log_stream: Optional[str] = None,
        log_readers: Optional[Sequence[PipesLogReader]] = None,
        max_retries: Optional[int] = DEFAULT_CLOUDWATCH_LOGS_MAX_RETRIES,
    ):
        """Args:
        client (boto3.client): boto3 CloudWatch client.
        """
        self.client: "CloudWatchLogsClient" = client or boto3.client("logs")
        self.log_group = log_group
        self.log_stream = log_stream
        self.max_retries = max_retries

        self.start_time = datetime.now()

        super().__init__(log_readers=log_readers)

    def on_launched(self, launched_payload: PipesLaunchedData) -> None:
        if "log_group" in launched_payload["extras"]:
            self.log_group = launched_payload["extras"]["log_group"]

        if "log_stream" in launched_payload["extras"]:
            self.log_stream = launched_payload["extras"]["log_stream"]

        self.launched_payload = launched_payload

    @contextmanager
    def get_params(self) -> Iterator[PipesParams]:
        yield {PipesDefaultMessageWriter.STDIO_KEY: PipesDefaultMessageWriter.STDOUT}

    def messages_are_readable(self, params: PipesParams) -> bool:
        if self.log_group is not None and self.log_stream is not None:
            # check if the stream actually exists
            try:
                resp = self.client.describe_log_streams(
                    logGroupName=self.log_group,
                    logStreamNamePrefix=self.log_stream,
                )

                if resp.get("logStreams", []):
                    return True
                else:
                    return False
            except self.client.exceptions.ResourceNotFoundException:
                return False
        else:
            return False

    def download_messages(
        self, cursor: Optional[str], params: PipesParams
    ) -> Optional[Tuple[str, str]]:
        params = {
            "logGroupName": self.log_group,
            "logStreamName": self.log_stream,
            "startTime": int(self.start_time.timestamp() * 1000),
        }

        if cursor is not None:
            params["nextToken"] = cursor

        response = get_log_events(client=self.client, max_retries=self.max_retries, **params)

        events = response.get("events")

        if not events:
            return None
        else:
            cursor = cast(str, response["nextForwardToken"])
            return cursor, "\n".join(
                cast(str, event.get("message")) for event in events if event.get("message")
            )

    def no_messages_debug_text(self) -> str:
        return "Attempted to read messages by extracting them from CloudWatch logs."
