HW-10-Socket-Programming/src/main/java/com/university/chat/Server/ClientSession.java

package com.university.chat.Server;

import com.university.chat.Common.ChatMessage;
import com.university.chat.Common.FileMessage;
import com.university.chat.Common.MessageType;

import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.Socket;
import java.nio.file.Files;

public class ClientSession implements Runnable {

    private String username;
    private Socket socket;
    private UserManager userManager;
    private ObjectInputStream in;
    private ObjectOutputStream out;

    public ClientSession(Socket socket, UserManager userManager) {
        try {
            this.socket = socket;
            this.userManager = userManager;

            out = new ObjectOutputStream(socket.getOutputStream());
            in = new ObjectInputStream(socket.getInputStream());

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void run() {
        try {
            Object firstObject = in.readObject();

            if (!(firstObject instanceof ChatMessage loginMessage)) {
                socket.close();
                return;
            }

            if (loginMessage.getType() != MessageType.LOGIN) {
                socket.close();
                return;
            }

            username = loginMessage.getSender();

            boolean added = userManager.addUser(username, this);

            if (!added) {
                ChatMessage failedMessage = new ChatMessage(
                        MessageType.LOGIN_FAILED,
                        "SERVER",
                        username,
                        "Username already exists"
                );

                send(failedMessage);
                socket.close();
                return;
            }

            FileManager.createUserFolders(username);

            ChatMessage successMessage = new ChatMessage(
                    MessageType.LOGIN_SUCCESS,
                    "SERVER",
                    username,
                    "Login successful"
            );

            send(successMessage);

            while (true) {
                Object object = in.readObject();

                if (object instanceof ChatMessage msg) {
                    handleChatMessage(msg);
                }

                if (object instanceof FileMessage fileMsg) {
                    handleFileMessage(fileMsg);
                }
            }

        } catch (Exception e) {
            System.out.println("Disconnected: " + username);
        } finally {
            if (username != null) {
                userManager.removeUser(username);
            }
        }
    }

    private void handleChatMessage(ChatMessage msg) throws IOException {
        switch (msg.getType()) {

            case PUBLIC_MESSAGE -> {
                for (ClientSession session : userManager.getAllSessions()) {
                    session.send(msg);
                }
            }

            case PRIVATE_MESSAGE -> {
                ClientSession receiver = userManager.getUser(msg.getReceiver());

                if (receiver != null) {
                    receiver.send(msg);
                }
            }

            case USER_LIST -> {
                ChatMessage userListMessage = new ChatMessage(
                        MessageType.USER_LIST,
                        "SERVER",
                        username,
                        userManager.listUsers()
                );

                send(userListMessage);
            }

            default -> {
            }
        }
    }

    private void handleFileMessage(FileMessage fileMsg) throws IOException {
        var sentPath = FileManager.getSentPath(
                fileMsg.getSender(),
                fileMsg.getFilename()
        );

        var recvPath = FileManager.getReceivedPath(
                fileMsg.getReceiver(),
                fileMsg.getFilename()
        );

        Files.write(sentPath, fileMsg.getData());
        Files.write(recvPath, fileMsg.getData());

        ClientSession receiver = userManager.getUser(fileMsg.getReceiver());

        if (receiver != null) {
            receiver.send(fileMsg);
        }
    }

    public void send(ChatMessage msg) throws IOException {
        out.writeObject(msg);
        out.flush();
    }

    public void send(FileMessage msg) throws IOException {
        out.writeObject(msg);
        out.flush();
    }
}