From 795f020a98cdc5a70b7953ffd4b981991e64ede6 Mon Sep 17 00:00:00 2001
From: Steven Tracey <steven@nevets.tech>
Date: Thu, 11 Jan 2024 08:48:11 -0600
Subject: [PATCH] Internal changes

---
 build.gradle                                  |  2 +-
 .../tech/nevets/dliteserver/Connection.java   | 16 ++++-
 .../dliteserver/MessageDistributer.java       | 61 +++++++++++++------
 3 files changed, 56 insertions(+), 23 deletions(-)

diff --git a/build.gradle b/build.gradle
index 563b86a..393c4c7 100644
--- a/build.gradle
+++ b/build.gradle
@@ -4,7 +4,7 @@ plugins {
 }
 
 group = 'tech.nevets'
-version = '1.0.0'
+version = '1.1.0'
 
 repositories {
     mavenCentral()
diff --git a/src/main/java/tech/nevets/dliteserver/Connection.java b/src/main/java/tech/nevets/dliteserver/Connection.java
index bdd44d8..55307f6 100644
--- a/src/main/java/tech/nevets/dliteserver/Connection.java
+++ b/src/main/java/tech/nevets/dliteserver/Connection.java
@@ -25,11 +25,23 @@ public class Connection implements Runnable {
     @Override
     public void run() {
         try {
-            md.registerStream(username, os);
+            md.registerStream(this);
             is.transferTo(md);
         } catch (IOException e) {
             System.out.println("Connection Disconnected");
         }
-        MessageDistributer.closeUserConnection(username);
+        MessageDistributer.closeUserConnection(this);
+    }
+
+    public String getUsername() {
+        return username;
+    }
+
+    public OutputStream getOutputStream() {
+        return this.os;
+    }
+
+    public void updateUsername(String username) {
+        this.username = username;
     }
 }
diff --git a/src/main/java/tech/nevets/dliteserver/MessageDistributer.java b/src/main/java/tech/nevets/dliteserver/MessageDistributer.java
index c07e6da..89b8033 100644
--- a/src/main/java/tech/nevets/dliteserver/MessageDistributer.java
+++ b/src/main/java/tech/nevets/dliteserver/MessageDistributer.java
@@ -3,15 +3,17 @@ package tech.nevets.dliteserver;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 public class MessageDistributer extends OutputStream {
-    private static Map<String, OutputStream> userStreams = new HashMap<>();
-
+    private static List<Connection> userStreams = new ArrayList<>();
+    private static MessageDistributer instance;
     private String buf = "";
 
+    public MessageDistributer() {
+        instance = this;
+    }
+
     @Override
     public void write(int b) throws IOException {
         buf += (char) b;
@@ -20,9 +22,9 @@ public class MessageDistributer extends OutputStream {
 
     @Override
     public void flush() {
-        for (OutputStream os : userStreams.values()) {
+        for (Connection conn : userStreams) {
             try {
-                os.write(buf.getBytes(StandardCharsets.UTF_8));
+                conn.getOutputStream().write((conn.getUsername() + "> " + buf).getBytes(StandardCharsets.UTF_8));
             } catch (IOException e) {
                 e.printStackTrace();
             }
@@ -30,21 +32,40 @@ public class MessageDistributer extends OutputStream {
         buf = "";
     }
 
-    public void registerStream(String username, OutputStream os) {
-        Set<String> usernames = userStreams.keySet();
-        while (usernames.contains(username)) {
-            username += ".imposter";
-        }
-
-        userStreams.put(username, os);
-        try {
-            this.write(("SERVER> Welcome, " + username + "!\n").getBytes(StandardCharsets.UTF_8));
-        } catch (IOException e) {
-            e.printStackTrace();
+    private void writeSystem(String message) {
+        for (Connection conn : userStreams) {
+            try {
+                conn.getOutputStream().write(("SERVER> " + message + "\n").getBytes(StandardCharsets.UTF_8));
+            } catch (IOException e) {
+                e.printStackTrace();
+            }
         }
     }
 
-    public static void closeUserConnection(String username) {
-        userStreams.remove(username);
+    public void registerStream(Connection conn) {
+        int i = 0;
+        String username;
+        for (username = conn.getUsername(); existingUsername(username); username = conn.getUsername() + "." + i) {
+            i++;
+        }
+        conn.updateUsername(username);
+        userStreams.add(conn);
+        writeSystem("Welcome " + username + "!");
+    }
+
+    private boolean existingUsername(String username) {
+        for (Connection c : userStreams) {
+            if (c.getUsername().equals(username)) return true;
+        }
+        return false;
+    }
+
+    private static MessageDistributer getInstance() {
+        return instance;
+    }
+
+    public static void closeUserConnection(Connection conn) {
+        userStreams.remove(conn);
+        getInstance().writeSystem(("Goodbye, " + conn.getUsername() + "!"));
     }
 }