import threading
import uuid
from typing import (
    Dict,
    Tuple,
)

import zmq


class LocalThreadData(threading.local):
    def __init__(self, context: zmq.Context):
        super(LocalThreadData, self).__init__()
        self.context = context

        self.addr_type_sockets: Dict[Tuple[str, zmq.SocketType], zmq.Socket] = {}
        self.socket_addr: Dict[zmq.Socket, str] = {}

    def open_socket(self, address, socket_type):
        key = (address, socket_type)
        if key in self.addr_type_sockets:
            return self.addr_type_sockets[key]

        socket = zmq.Socket(self.context, socket_type)
        socket.set_string(zmq.IDENTITY, str(uuid.uuid4()))
        socket.connect(addr=address)
        self.addr_type_sockets[key] = socket
        self.socket_addr[socket] = address
        return socket

    def close_socket(self, socket: zmq.Socket):
        if socket in self.socket_addr:
            key = (self.socket_addr[socket], socket.type)
            self.addr_type_sockets.pop(key, None)
            del self.socket_addr[socket]

        socket.close(linger=0)

    def close_all_sockets(self):
        for socket in self.socket_addr.keys():
            socket.close()

    def __del__(self):
        self.close_all_sockets()
