diff --git a/src/mind_reader/core/__init__.py b/src/mind_reader/core/__init__.py
index b5e32ec..6bbc815 100644
--- a/src/mind_reader/core/__init__.py
+++ b/src/mind_reader/core/__init__.py
@@ -1,5 +1,7 @@
-from mind_reader.core.game_state import GameState
+from .state import AppState, LobbyState, SessionState
 
 __all__ = [
-    "GameState",
+    "AppState",
+    "LobbyState",
+    "SessionState",
 ]
diff --git a/src/mind_reader/core/game_state.py b/src/mind_reader/core/game_state.py
deleted file mode 100644
index 566463e..0000000
--- a/src/mind_reader/core/game_state.py
+++ /dev/null
@@ -1,33 +0,0 @@
-class GameState:
-    def __init__(self, local_player_name: str):
-        self.local_player_name: str = local_player_name
-        self.remote_player_name: str | None = None
-        self.lobby_players: dict[str, str] = {self.local_player_name: "NONE"}
-
-    def add_remote_player(self, player_name: str):
-        self.remote_player_name = player_name
-        self.lobby_players[self.remote_player_name] = "NONE"
-
-    def delete_remote_player(self):
-        if self.remote_player_name in self.lobby_players:
-            del self.lobby_players[self.remote_player_name]
-        self.remote_player_name = None
-
-    def set_player_role(self, player_name: str, role: str) -> bool:
-        current_role = self.lobby_players.get(player_name)
-        if current_role == role:
-            return False
-
-        if role in ["NONE", "SETTER", "GUESSER"]:
-            self.lobby_players[player_name] = role
-            return True
-
-        return False
-
-    @property
-    def local_player_role(self) -> str:
-        return self.lobby_players[self.local_player_name]
-
-    def is_session_ready(self) -> bool:
-        roles = list(self.lobby_players.values())
-        return ("SETTER" in roles) and ("GUESSER" in roles)
diff --git a/src/mind_reader/core/state.py b/src/mind_reader/core/state.py
new file mode 100644
index 0000000..b89de3d
--- /dev/null
+++ b/src/mind_reader/core/state.py
@@ -0,0 +1,172 @@
+from typing import Callable, Optional
+
+from wx.lib.pubsub import pub
+
+from ..network import NetworkController, NetworkMesssage
+
+OnNetworkEventCallback = Callable[[dict[str, str]], None]
+
+
+class BaseState:
+    def __init__(self, app_state: "AppState"):
+        self.app_state = app_state
+
+    def handle_network_message(self, message: NetworkMesssage) -> None:
+        pass
+
+
+class LobbyState(BaseState):
+    def __init__(self, app_state: "AppState", local_player_name: str):
+        super().__init__(app_state)
+        self.local_player_name: str = local_player_name
+        self.remote_player_name: Optional[str] = None
+        self.players: dict[str, str] = {self.local_player_name: "NONE"}
+
+    def handle_network_message(self, message: NetworkMesssage) -> None:
+        callbacks: dict[str, OnNetworkEventCallback] = {
+            "add_player": self.handle_add_player,
+            "get_player": self.handle_get_player,
+            "set_role": self.handle_set_player_role,
+            "disconnect": self.handle_disconnect,
+            "server_stop": self.handle_server_stop,
+        }
+
+        callback = callbacks.get(message.event, None)
+        if callback:
+            callback(message.payload)
+            pub.sendMessage("lobby.refresh", state=self)
+
+    def send_local_player(self) -> None:
+        player_name = self.local_player_name
+        player_role = self.players.get(player_name, "NONE")
+
+        self.app_state.network.send(
+            "add_player",
+            {
+                "player_name": player_name,
+                "player_role": player_role,
+            },
+        )
+
+    def get_remote_player(self) -> None:
+        self.app_state.network.send("get_player")
+
+    def set_player_role(self, role: str) -> bool:
+        current_role = self.players[self.local_player_name]
+        result = False
+
+        if current_role == role:
+            result = False
+
+        if role in ["NONE", "SETTER", "GUESSER"]:
+            self.players[self.local_player_name] = role
+            result = True
+
+        if result:
+            if self.app_state.network.conn:
+                self.app_state.network.send(
+                    "set_role",
+                    {
+                        "player_name": self.local_player_name,
+                        "player_role": role,
+                    },
+                )
+
+        return result
+
+    def send_client_disconnect(self) -> None:
+        self.app_state.network.send(
+            "disconnect", {"player_name": self.local_player_name}
+        )
+
+    def send_server_stop(self) -> None:
+        self.app_state.network.send("server_stop")
+
+    def handle_add_player(self, payload: dict[str, str]) -> None:
+        player_name = payload["player_name"]
+        player_role = payload["player_role"]
+        self.remote_player_name = player_name
+        self.players[self.remote_player_name] = player_role
+
+    def handle_get_player(self, payload: dict[str, str]) -> None:
+        self.send_local_player()
+
+    def handle_set_player_role(self, payload: dict[str, str]) -> None:
+        player_name = payload["player_name"]
+        player_role = payload["player_role"]
+
+        self.players[player_name] = player_role
+
+    def handle_disconnect(self, payload: dict[str, str]) -> None:
+        player_name = payload["player_name"]
+        self.players.pop(player_name)
+        self.remote_player_name = None
+
+    def handle_server_stop(self, payload: dict[str, str]) -> None:
+        pub.sendMessage("mainframe.remote_server_stop", state=self.app_state)
+
+    @property
+    def local_player_role(self) -> str:
+        return self.players[self.local_player_name]
+
+    def is_ready(self) -> bool:
+        roles = list(self.players.values())
+        return ("SETTER" in roles) and ("GUESSER" in roles)
+
+
+class SessionState(BaseState):
+    def __init__(self, app_state: "AppState"):
+        super().__init__(app_state)
+
+
+class AppState:
+    def __init__(self):
+        self.network = NetworkController(
+            self.handle_network_message,
+            self.handle_network_notice,
+        )
+
+        self.lobby: Optional[LobbyState] = None
+        self.session: Optional[SessionState] = None
+
+    def handle_network_message(self, message: NetworkMesssage) -> None:
+        if self.lobby:
+            self.lobby.handle_network_message(message)
+
+        if self.session:
+            self.session.handle_network_message(message)
+
+    def handle_network_notice(self, notice: str) -> None:
+        pub.sendMessage("mainframe.statusbar", notice=notice)
+
+    def recreate_network_controller(self) -> None:
+        self.network.stop()
+        self.network = NetworkController(
+            self.handle_network_message,
+            self.handle_network_notice,
+        )
+
+    def connect(self, player_name: str, host: str, port: int = 8994) -> bool:
+        if self.network.connect(host, port):
+            self.lobby = LobbyState(self, player_name)
+            self.lobby.send_local_player()
+            self.lobby.get_remote_player()
+            return True
+        else:
+            self.recreate_network_controller()
+            return False
+
+    def serve(self, player_name: str, host: str, port: int = 8994) -> None:
+        self.network.serve(host, port)
+        self.lobby = LobbyState(self, player_name)
+
+    def disconnect(self) -> None:
+        if self.lobby:
+            if self.network.is_server:
+                self.lobby.send_server_stop()
+            else:
+                self.lobby.send_client_disconnect()
+
+            self.lobby = None
+
+        self.recreate_network_controller()