from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any, Dict, Iterator, AsyncIterator, cast

from ..._utils import lru_cache
from ..._streaming import ServerSentEvent

if TYPE_CHECKING:
    from botocore.model import Shape
    from botocore.eventstream import EventStreamMessage


@lru_cache(maxsize=None)
def get_response_stream_shape() -> Shape:
    from botocore.model import ServiceModel
    from botocore.loaders import Loader

    loader = Loader()
    bedrock_service_dict = loader.load_service_model("bedrock-runtime", "service-2")
    bedrock_service_model = ServiceModel(bedrock_service_dict)
    return bedrock_service_model.shape_for("ResponseStream")


class AWSEventStreamDecoder:
    def __init__(self) -> None:
        from botocore.parsers import EventStreamJSONParser

        self.parser = EventStreamJSONParser()

    def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[ServerSentEvent]:
        """Given an iterator that yields lines, iterate over it & yield every event encountered"""
        from botocore.eventstream import EventStreamBuffer

        event_stream_buffer = EventStreamBuffer()
        for chunk in iterator:
            event_stream_buffer.add_data(chunk)
            for event in event_stream_buffer:
                sse = self._parse_message_from_event(event)
                if sse:
                    yield sse

    async def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]:
        """Given an async iterator that yields lines, iterate over it & yield every event encountered"""
        from botocore.eventstream import EventStreamBuffer

        event_stream_buffer = EventStreamBuffer()
        async for chunk in iterator:
            event_stream_buffer.add_data(chunk)
            for event in event_stream_buffer:
                sse = self._parse_message_from_event(event)
                if sse:
                    yield sse

    def _parse_message_from_event(self, event: EventStreamMessage) -> ServerSentEvent | None:
        response_dict = event.to_response_dict()
        parsed_response = self.parser.parse(response_dict, get_response_stream_shape())
        if response_dict["status_code"] != 200:
            raise ValueError(f"Bad response code, expected 200: {response_dict}")

        chunk = parsed_response.get("chunk")
        if not chunk:
            return None

        return _chunk_bytes_to_sse(chunk.get("bytes"))


def _chunk_bytes_to_sse(raw: bytes) -> ServerSentEvent | None:
    decoded = raw.decode()
    data: Any
    try:
        data = json.loads(decoded)
    except Exception:
        data = None

    if not isinstance(data, dict):
        return ServerSentEvent(data=decoded, event="completion")

    payload = cast("Dict[str, Any]", data)
    event_type = payload.get("type")
    if not isinstance(event_type, str):
        event_type = "completion"

    return ServerSentEvent(data=decoded, event=event_type)
