Assignment-10/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 final Socket socket;
    private final UserManager userManager;
    private ObjectOutputStream out;
    private ObjectInputStream in;
    private String username;

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

        try {
            this.out = new ObjectOutputStream(socket.getOutputStream());
            this.out.flush();
            this.in = new ObjectInputStream(socket.getInputStream());
        } catch (IOException e) {
            System.err.println("Error initializing streams for session: " + e.getMessage());
        }
    }

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

            if (firstObject instanceof ChatMessage loginMsg && loginMsg.getType() == MessageType.LOGIN) {
                this.username = loginMsg.getSender();

                if (userManager.addUser(username, this)) {
                    FileManager.createUserFolders(username);
                    sendMessage(new ChatMessage(MessageType.LOGIN_SUCCESS, "SERVER", username, "Login successful."));
                    System.out.println("[Server] User logged in: " + username);
                } else {
                    sendMessage(new ChatMessage(MessageType.LOGIN_FAILED, "SERVER", username, "Username is already taken."));
                    socket.close();
                    return;
                }
            } else {
                socket.close();
                return;
            }

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

                if (receivedObject instanceof ChatMessage) {
                    handleChatMessage((ChatMessage) receivedObject);
                } else if (receivedObject instanceof FileMessage) {
                    handleFileMessage((FileMessage) receivedObject);
                }
            }
        } catch (Exception e) {
            System.out.println("Disconnected: " + username);
        } finally {
            if (username != null) {
                if (userManager.getUser(username) == this) {
                    userManager.removeUser(username);
                    System.out.println("[Server] Safely removed active user: " + username);
                }
            }
            try {
                socket.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private void handleChatMessage(ChatMessage msg) throws IOException {
        switch (msg.getType()) {
            case PUBLIC_MESSAGE -> {
                for (ClientSession session : userManager.getAllSessions()) {
                    session.sendMessage(msg);
                }
            }
            case PRIVATE_MESSAGE -> {
                ClientSession recipient = userManager.getUser(msg.getReceiver());
                if (recipient != null) {
                    recipient.sendMessage(msg);
                } else {
                    sendMessage(new ChatMessage(MessageType.PUBLIC_MESSAGE, "SERVER", username, "User " + msg.getReceiver() + " is offline."));
                }
            }
            case USER_LIST -> {
                String activeUsers = userManager.listUsers();
                ChatMessage listMsg = new ChatMessage(MessageType.USER_LIST, "SERVER", username, activeUsers);
                sendMessage(listMsg);
            }
            default -> System.out.println("Unknown ChatMessage type received.");
        }
    }

    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 recipient = userManager.getUser(fileMsg.getReceiver());
        if (recipient != null) {
            recipient.sendMessage(fileMsg);
        }
    }

    public synchronized void sendMessage(Object msg) {
        try {
            if (out != null) {
                out.writeObject(msg);
                out.flush();
            }
        } catch (IOException e) {
            System.err.println("Failed to send packet to " + username + ": " + e.getMessage());
        }
    }
}