import ssl
import socket
import threading
import struct
from . import Socket, NetworkError


class TCPReq(Socket):
    def __init__(self, port):
        self.conn = None
        self._lock = threading.RLock()
        self.port = port
        self._retries = 0

    def connect(self, ip_addr):
        self.ip_addr = ip_addr

        context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
        context.check_hostname = False
        context.verify_mode = ssl.CERT_NONE

        family, sock_type, proto, _, sockaddr = socket.getaddrinfo(
            host=ip_addr, port=self.port, proto=socket.IPPROTO_TCP)[0]

        sock = socket.socket(family, sock_type, proto)
        conn = context.wrap_socket(sock)
        conn.connect(sockaddr)

        self.sock = sock
        self.conn = conn

    def transfer(self, data, timeout=1000):
        # TODO use timeout
        length = struct.pack('<I', len(data))
        with self._lock:
            try:
                self.conn.sendall(length + data)
                length = struct.unpack('<I', self.conn.recv(4))[0]
                self._retries = 0
            except (BrokenPipeError, struct.error, OSError):
                # recursive reconnect
                self.close()

                if self._retries > 5:
                    raise
                self._retries += 1

                self.connect(self.ip_addr)
                return self.transfer(data)

            rep = b''
            while len(rep) < length:
                rep += self.conn.recv(length - len(rep))

        ptype, status = struct.unpack("<BB", rep[:2])

        # Response should begin with the same packet type, and either 0 or
        # <5 status code
        if data[0] != ptype or (status <= 5 and status > 0):
            raise NetworkError("Status: 0x{:02X}, {:s}".format(
                status, rep[2:].decode('utf-8')))
        return rep

    def close(self):
        try:
            with self._lock:
                self.conn.close()
                self.sock.close()
        except Exception:
            pass


class TCPRenderReq(TCPReq):
    def transfer(self, data):
        resp = super().transfer(b'\x57' + data)
        assert resp[:2] == b'\x57\x00'
        return resp[2:]


class TCPSocketFactory:
    def __init__(self, ip_addr):
        self.ip_addr = ip_addr

    def get_control_socket(self):
        skt = TCPReq(27181)
        skt.connect(self.ip_addr)
        return skt

    def get_render_socket(self):
        skt = TCPRenderReq(27181)
        skt.connect(self.ip_addr)
        return skt
