From f435d16276f35fd3b95bf8bd0f2e7d957ac983fb Mon Sep 17 00:00:00 2001
From: Brikwerk <reece@brikwerk.com>
Date: Sun, 26 Sep 2021 21:15:31 -0700
Subject: [PATCH] Fixed reconnection issues

---
 nxbt/controller/server.py | 46 ++++++++++++++++++++++++++++++++++++---
 nxbt/nxbt.py              |  5 +++++
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/nxbt/controller/server.py b/nxbt/controller/server.py
index 2cbd8e1..633e1ad 100644
--- a/nxbt/controller/server.py
+++ b/nxbt/controller/server.py
@@ -98,7 +98,8 @@ class ControllerServer():
                 if self.lock:
                     self.lock.release()
 
-            self.switch_address = itr.getsockname()[0]
+            self.switch_address = itr.getpeername()[0]
+            self.state["last_connection"] = self.switch_address
 
             self.state["state"] = "connected"
 
@@ -166,10 +167,9 @@ class ControllerServer():
                     self.cached_msg = msg[3:]
                 # Send a blank packet every so often to keep the Switch
                 # from disconnecting from the controller.
-                elif self.tick >= 1320:
+                elif self.tick >= 132:
                     itr.sendall(msg)
                     self.tick = 0
-                    # print(msg, "tick")
             except BlockingIOError:
                 continue
             except OSError as e:
@@ -207,10 +207,50 @@ class ControllerServer():
                     self.bt.address,
                     colour_body=self.colour_body,
                     colour_buttons=self.colour_buttons)
+                self.input.reassign_protocol(self.protocol)
                 if self.lock:
                     self.lock.acquire()
                 try:
                     itr, ctrl = self.reconnect(self.switch_address)
+
+                    received_first_message = False
+                    while True:
+                        # Attempt to get output from Switch
+                        try:
+                            reply = itr.recv(50)
+                            if self.logger_level <= logging.DEBUG and len(reply) > 40:
+                                self.logger.debug(format_msg_switch(reply))
+                        except BlockingIOError:
+                            reply = None
+
+                        if reply:
+                            received_first_message = True
+
+                        self.protocol.process_commands(reply)
+                        msg = self.protocol.get_report()
+
+                        if self.logger_level <= logging.DEBUG and reply:
+                            self.logger.debug(format_msg_controller(msg))
+
+                        try:
+                            itr.sendall(msg)
+                        except BlockingIOError:
+                            continue
+
+                        # Exit pairing loop when player lights have been set and
+                        # vibration has been enabled
+                        if (reply and len(reply) > 45 and
+                                self.protocol.vibration_enabled and self.protocol.player_number):
+                            break
+
+                        # Switch responds to packets slower during pairing
+                        # Pairing cycle responds optimally on a 15Hz loop
+                        if not received_first_message:
+                            time.sleep(1)
+                        else:
+                            time.sleep(1/15)
+
+                    self.state["state"] = "connected"
                     return itr, ctrl
                 finally:
                     if self.lock:
diff --git a/nxbt/nxbt.py b/nxbt/nxbt.py
index 3ae411f..780be29 100644
--- a/nxbt/nxbt.py
+++ b/nxbt/nxbt.py
@@ -740,6 +740,11 @@ class _ControllerManager():
         controller_state["finished_macros"] = []
         controller_state["errors"] = False
         controller_state["direct_input"] = json.loads(json.dumps(DIRECT_INPUT_PACKET))
+        controller_state["colour_body"] = colour_body
+        controller_state["colour_buttons"] = colour_buttons
+        controller_state["type"] = str(controller_type)
+        controller_state["adapter_path"] = adapter_path
+        controller_state["last_connection"] = None
 
         self._controller_queues[index] = controller_queue