import socket
import struct
import time
import numpy as np


class TcpService(socket.socket):
    _init_timeout = 1  # 连接/发送使用的超时时间, 接收使用总时间-发送耗时
    _block_size = 4096

    @classmethod
    def get_service(cls, ip, port, timeout=None, ip_type=socket.AF_INET, sock_type=socket.SOCK_STREAM):
        try:
            self = cls(ip_type, sock_type)
            self.settimeout(timeout)
            st = time.time()
            self.connect((ip, port))
            print(time.time() - st)
            setattr(self, "_total_timeout", timeout - (time.time() - st))
            self.settimeout(self._init_timeout)
            return self
        except ConnectionRefusedError as e:
            return f"{e}"
        except Exception as e:
            return f"{e}"

    @classmethod
    def online(cls, address, timeout=1, ip_type=socket.AF_INET, sock_type=socket.SOCK_STREAM):
        self = cls(ip_type, sock_type)
        self.settimeout(timeout)
        status_code = self.connect_ex(address)
        self.close()
        return status_code

    def send_data(self, data: iter, fmt: str):
        data = struct.pack(fmt, *data)
        total_length, sent_length = len(data), 0
        st = time.time()
        while total_length != sent_length:
            sent_length += self.send(data[sent_length: self._block_size + sent_length])
            assert time.time() - st <= self._total_timeout, f"socket.send total timeout {self._total_timeout}, sent {sent_length}"

        self._total_timeout -= (time.time() - st)
        print(self._total_timeout)
        self.settimeout(self._total_timeout)

    def recv_data(self, bufsize: int, fmt: str):
        try:
            block_size, bytes_data, bytes_data_length = self._block_size, b"", 0
            st = time.time()
            while bytes_data_length != bufsize:
                if not (bufsize - bytes_data_length) // self._block_size:
                    block_size = (bufsize - bytes_data_length) % self._block_size
                bytes_data += self.recv(block_size)
                bytes_data_length = len(bytes_data)
                gap_time = time.time() - st
                assert gap_time <= self._total_timeout, f"left timeout {self._total_timeout}, recv {bytes_data_length}"
                self.settimeout(self._total_timeout - gap_time)
        except (socket.timeout, AssertionError) as e:
            assert 0, f"socket.recv {e}"
        return True, struct.unpack(fmt, bytes_data)


if __name__ == '__main__':
    length = 1024 ** 2 * 4
    fmt = "=" + "I" * length
    data = np.arange(length, dtype='u4')
    socket_client = TcpService.get_service("192.168.1.65", 8001, timeout=10)
    assert not isinstance(socket_client, str), socket_client
    with socket_client as s:
        s.send_data(data, fmt)