import json
from socket import socket
from typing import AnyStr


def _decode(data: AnyStr, decode: bool):
    if decode:
        data = data.decode("utf-8")
    return data


def _encode(data: AnyStr, encode: bool):
    if encode:
        data = data.encode("utf-8")
    return data


class Buff(object):
    def __init__(self, sock: socket):
        """

        :param sock:
        """
        self.wbuff = sock.makefile("wb")
        self.rbuff = sock.makefile("rb")

    def read(self, size: int, decode: bool = False) -> AnyStr:
        return _decode(self.rbuff.read(size), decode)

    def readLine(self, deocde: bool = False) -> AnyStr:
        return _decode(self.rbuff.readline(), deocde)

    def write(self, data: AnyStr, encdoe: bool = True):
        self.wbuff.write(_encode(data, encdoe))

    def writeLine(self, data: AnyStr, encdoe: bool = True):
        if isinstance(data, bytes):
            data += b"\r\n"
        else:
            data += "\r\n"

        self.write(data, encdoe)

    def writeJson(self, data: dict, encode: bool = True, wrap: bool = False):
        data = json.dumps(data)
        data = _encode(data, encode)
        if wrap:
            self.writeLine(data, False)
        self.write(data, False)

    def send(self):
        self.wbuff.flush()

    def close(self):
        self.wbuff.close()
        self.rbuff.close()