import sys
import socket
import datetime

import anyio
from async_generator import async_generator, yield_, yield_from_
import h2
import h2.events
import h2.exceptions
from purerpc.utils import is_darwin
from purerpc.grpclib.exceptions import ProtocolError

from .grpclib.connection import GRPCConfiguration, GRPCConnection
from .grpclib.events import RequestReceived, RequestEnded, ResponseEnded, MessageReceived
from .grpclib.buffers import MessageWriteBuffer, MessageReadBuffer


class SocketWrapper:
    def __init__(self, grpc_connection: GRPCConnection, sock: anyio._networking.SocketStream):
        self._set_raw_socket_options(sock._socket._raw_socket)
        self._socket = sock._socket
        self._grpc_connection = grpc_connection
        self._write_fifo_lock = anyio.create_lock()
        self._retry = False

    @staticmethod
    def _set_raw_socket_options(raw_socket):
        raw_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
        if hasattr(socket, "TCP_KEEPIDLE"):
            raw_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 300)
        elif is_darwin():
            # Darwin specific option
            TCP_KEEPALIVE = 16
            raw_socket.setsockopt(socket.IPPROTO_TCP, TCP_KEEPALIVE, 300)
        raw_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 30)
        raw_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5)
        raw_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

    async def flush(self):
        """This maybe called from different threads."""
        async with self._write_fifo_lock:
            self._retry = True
            while self._retry:
                self._retry = False
                data = self._grpc_connection.data_to_send()
                if not data:
                    return
                await self._socket.sendall(data)

    async def try_flush(self):
        if self.locked:
            self._retry = True
        else:
            await self.flush()

    @property
    def locked(self):
        return self._write_fifo_lock.locked()

    async def recv(self, buffer_size: int):
        """This may only be called from single thread."""
        return await self._socket.recv(buffer_size)


class GRPCStream:
    def __init__(self, grpc_connection: GRPCConnection, stream_id: int, socket: SocketWrapper,
                 grpc_socket: "GRPCSocket"):
        self._stream_id = stream_id
        self._grpc_connection = grpc_connection
        self._grpc_socket = grpc_socket
        self._socket = socket
        self._flow_control_update_event = anyio.create_event()
        self._incoming_events = anyio.create_queue(sys.maxsize)

        self._response_started = False
        self._closed = False
        
        self._start_stream_event = None
        self._end_stream_event = None

    @property
    def start_stream_event(self):
        return self._start_stream_event

    @property
    def end_stream_event(self):
        return self._end_stream_event

    @property
    def stream_id(self):
        return self._stream_id

    @property
    def client_side(self):
        return self._grpc_connection.config.client_side

    @property
    def debug_prefix(self):
        return "[CLIENT] " if self.client_side else "[SERVER] "

    async def _set_flow_control_update(self):
        await self._flow_control_update_event.set()

    async def _wait_flow_control_update(self):
        await self._flow_control_update_event.wait()
        self._flow_control_update_event.clear()

    async def _send(self, message: bytes, compress=False):
        message_write_buffer = MessageWriteBuffer(self._grpc_connection.config.message_encoding,
                                                  self._grpc_connection.config.max_message_length)
        message_write_buffer.write_message(message, compress)
        while message_write_buffer:
            window_size = self._grpc_connection.flow_control_window(self._stream_id)
            if window_size <= 0:
                await self._wait_flow_control_update()
                continue
            num_data_to_send = min(window_size, len(message_write_buffer))
            data = message_write_buffer.data_to_send(num_data_to_send)
            self._grpc_connection.send_data(self._stream_id, data)
            await self._socket.flush()

    async def _receive(self):
        event = await self._incoming_events.get()
        if isinstance(event, MessageReceived):
            self._grpc_connection.acknowledge_received_data(self._stream_id,
                                                            event.flow_controlled_length)
            await self._socket.flush()
        elif isinstance(event, RequestEnded) or isinstance(event, ResponseEnded):
            assert self._end_stream_event is None
            self._end_stream_event = event
        else:
            assert self._start_stream_event is None
            self._start_stream_event = event
        return event

    async def close(self, status=None, content_type_suffix="", custom_metadata=()):
        if self.client_side and (status or custom_metadata):
            raise ValueError("Client side streams cannot be closed with non-default arguments")
        if self._closed:
            raise TypeError("Closing already closed stream")
        self._closed = True
        if self.client_side:
            try:
                self._grpc_connection.end_request(self._stream_id)
            except h2.exceptions.StreamClosedError:
                # Remote end already closed connection, do nothing here
                pass
        elif self._response_started:
            self._grpc_connection.end_response(self._stream_id, status, custom_metadata)
        else:
            self._grpc_connection.respond_status(self._stream_id, status,
                                                 content_type_suffix, custom_metadata)
        await self._socket.flush()
        self._grpc_socket._decref_stream(self._stream_id)

    async def start_response(self, content_type_suffix="", custom_metadata=()):
        if self.client_side:
            raise ValueError("Cannot start response on client-side socket")
        self._grpc_connection.start_response(self._stream_id, content_type_suffix, custom_metadata)
        self._response_started = True
        await self._socket.flush()


# TODO: this name is not correct, should be something like GRPCConnection (but this name is already
# occupied)
class GRPCSocket:
    StreamClass = GRPCStream

    def __init__(self, config: GRPCConfiguration, sock,
                 receive_buffer_size=16384):
        self._grpc_connection = GRPCConnection(config=config)
        self._socket = SocketWrapper(self._grpc_connection, sock)
        self._receive_buffer_size = receive_buffer_size

        self._streams = {}  # type: Dict[int, GRPCStream]
        self._streams_count = {}  # type: Dict[int, int]

    @property
    def client_side(self):
        return self._grpc_connection.config.client_side

    def _stream_ctor(self, stream_id):
        return self.StreamClass(self._grpc_connection, stream_id, self._socket, self)

    def _allocate_stream(self, stream_id):
        stream = self._stream_ctor(stream_id)
        self._streams[stream_id] = stream
        self._streams_count[stream_id] = 2
        return stream

    def _decref_stream(self, stream_id: int):
        self._streams_count[stream_id] -= 1
        if self._streams_count[stream_id] == 0:
            del self._streams[stream_id]
            del self._streams_count[stream_id]

    @async_generator
    async def _listen(self):
        while True:
            data = await self._socket.recv(self._receive_buffer_size)
            if not data:
                return
            events = self._grpc_connection.receive_data(data)
            for event in events:
                # TODO: implement this
                if isinstance(event, h2.events.WindowUpdated):
                    if event.stream_id == 0:
                        for stream in self._streams.values():
                            await stream._set_flow_control_update()
                    elif event.stream_id in self._streams:
                        await self._streams[event.stream_id]._set_flow_control_update()
                    continue
                elif isinstance(event, RequestReceived):
                    self._allocate_stream(event.stream_id)

                await self._streams[event.stream_id]._incoming_events.put(event)

                if isinstance(event, RequestReceived):
                    await yield_(self._streams[event.stream_id])
                elif isinstance(event, ResponseEnded) or isinstance(event, RequestEnded):
                    self._decref_stream(event.stream_id)
            await self._socket.try_flush()

    async def _listener_thread(self):
        async for _ in self._listen():
            raise ProtocolError("Received request on client end")

    async def initiate_connection(self, parent_task_group: anyio.abc.TaskGroup):
        self._grpc_connection.initiate_connection()
        await self._socket.flush()
        if self.client_side:
            await parent_task_group.spawn(self._listener_thread)

    @async_generator
    async def listen(self):
        if self.client_side:
            raise ValueError("Cannot listen client-side socket")
        await yield_from_(self._listen())

    async def start_request(self, scheme: str, service_name: str, method_name: str,
                            message_type=None, authority=None, timeout: datetime.timedelta=None,
                            content_type_suffix="", custom_metadata=()):
        if not self.client_side:
            raise ValueError("Cannot start request on server-side socket")
        stream_id = self._grpc_connection.get_next_available_stream_id()
        stream = self._allocate_stream(stream_id)
        self._grpc_connection.start_request(stream_id, scheme, service_name, method_name,
                                            message_type, authority, timeout,
                                            content_type_suffix, custom_metadata)
        await self._socket.flush()
        return stream
