from __future__ import annotations

import sys
from typing import Callable, Optional

if sys.version_info >= (3, 8):
    from typing import Protocol
else:
    from typing_extensions import Protocol

from ..thrift import TType, TException


def readall(read_fn: Callable[[int], bytes], sz: int) -> bytes:
    buff = b''
    have = 0
    while have < sz:
        chunk = read_fn(sz - have)
        have += len(chunk)
        buff += chunk

        if len(chunk) == 0:
            raise TTransportException(TTransportException.END_OF_FILE,
                                      "End of file reading from transport")

    return buff


class TTransportFactory(Protocol):
    """Transport factory interface for type annotations."""

    def get_transport(self, trans: TTransportBase) -> TTransportBase:
        """Return a transport instance wrapping the given transport."""
        ...


class TTransportBase(object):
    """Base class for Thrift transport layer."""

    def is_open(self) -> bool:
        """Check if this transport is open."""
        raise NotImplementedError

    def open(self) -> None:
        """
        Prepare this transport for usage and allocate any necessary resources
        like sockets or sessions.
        """
        raise NotImplementedError

    def close(self) -> None:
        """Clean up and deallocate any resources allocated in open()."""
        raise NotImplementedError

    def _read(self, sz: int) -> bytes:
        """
        Internal read method which can read up to `sz` bytes but doesn't
        need to return them all.
        """
        raise NotImplementedError

    def read(self, sz: int) -> bytes:
        """
        Get exactly `sz` bytes from the underlying connection.

        When implementing a custom transport, this method must return exactly
        `sz` bytes if it is expected to be called from the protocol layer. If
        it intends to wrapped by another transport, like TBufferedTransport,
        it should return whatever the underlying connection/transport can get.
        The wrapping transport will take care of ensuring `sz` bytes are
        returned. For a more in depth discussion, see:
        https://github.com/Thriftpy/thriftpy2/pull/108#discussion_r355131677
        """
        return readall(self._read, sz)

    def write(self, buf: bytes) -> None:
        """
        Submit some data to be written to the connection. May be
        buffered until flush is called.
        """
        raise NotImplementedError

    def flush(self) -> None:
        """Ensure that all internal buffers are emptied into the connection."""
        raise NotImplementedError


class TTransportException(TException):
    """Custom Transport Exception class"""

    thrift_spec = {
        1: (TType.STRING, 'message'),
        2: (TType.I32, 'type'),
    }

    UNKNOWN = 0
    NOT_OPEN = 1
    ALREADY_OPEN = 2
    TIMED_OUT = 3
    END_OF_FILE = 4

    def __init__(self, type: int = UNKNOWN,
                 message: Optional[str] = None) -> None:
        super(TTransportException, self).__init__()
        self.type = type
        self.message = message
