RITz-WhatsApp-Bot/src/main/java/com/rit/whatsapp/service/SessionService.java

package com.rit.whatsapp.service;

import com.rit.whatsapp.model.UserEntity;
import com.rit.whatsapp.model.UserState;
import com.rit.whatsapp.model.GrievanceDraft;
import com.rit.whatsapp.repository.UserRepository;
import com.rit.whatsapp.util.EncryptionUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class SessionService {

    private final Map<String, UserState> userStates = new ConcurrentHashMap<>();
    
    // In-memory cache for API session tokens
    private final Map<String, String> sessionTokens = new ConcurrentHashMap<>();
    private final Map<String, String> temporaryLoginIds = new ConcurrentHashMap<>();
    private final Map<String, String> studentNames = new ConcurrentHashMap<>();

    private final UserRepository userRepository;
    private final EncryptionUtils encryptionUtils;
    private final ImsApiService imsApiService;

    @Autowired
    public SessionService(UserRepository userRepository, 
                          EncryptionUtils encryptionUtils,
                          ImsApiService imsApiService) {
        this.userRepository = userRepository;
        this.encryptionUtils = encryptionUtils;
        this.imsApiService = imsApiService;
    }

    public UserState getUserState(String phoneNumber) {
        return userStates.getOrDefault(phoneNumber, UserState.MAIN_MENU);
    }

    public void setUserState(String phoneNumber, UserState state) {
        userStates.put(phoneNumber, state);
    }
    
    public boolean isSessionCached(String phoneNumber) {
        return sessionTokens.containsKey(phoneNumber);
    }

    public boolean hasSavedCredentials(String phoneNumber) {
        return userRepository.existsByPhoneNumber(phoneNumber);
    }

    public void resetState(String phoneNumber) {
        userStates.remove(phoneNumber);
    }

    public boolean isAuthenticated(String phoneNumber) {
        if (sessionTokens.containsKey(phoneNumber)) {
            return true;
        }

        // Check database for persistent session
        Optional<UserEntity> userOpt = userRepository.findByPhoneNumber(phoneNumber);
        if (userOpt.isPresent()) {
            UserEntity user = userOpt.get();
            if (user.getEncryptedPasscode() == null) {
                return false;
            }
            String decryptedPasscode = encryptionUtils.decrypt(user.getEncryptedPasscode());
            
            // Silent background login
            ImsApiService.AuthResult result = imsApiService.authenticate(user.getImsId(), decryptedPasscode);
            if (result != null) {
                sessionTokens.put(phoneNumber, result.token);
                if (user.isFaculty() != result.isFaculty) {
                    user.setFaculty(result.isFaculty);
                    userRepository.save(user);
                }
                return true;
            }
        }
        
        return false;
    }

    public void authenticateUser(String phoneNumber, String imsId, String passcode, ImsApiService.AuthResult result) {
        sessionTokens.put(phoneNumber, result.token);
        
        // Persist to database (students table)
        Optional<UserEntity> userOpt = userRepository.findById(imsId);
        UserEntity user;
        if (userOpt.isPresent()) {
            user = userOpt.get();
            user.setPhoneNumber(phoneNumber);
            user.setEncryptedPasscode(encryptionUtils.encrypt(passcode));
        } else {
            String encryptedPasscode = encryptionUtils.encrypt(passcode);
            user = new UserEntity(imsId, phoneNumber, encryptedPasscode);
        }
        user.setFaculty(result.isFaculty);
        userRepository.save(user);
    }

    public String getImsToken(String phoneNumber) {
        // Ensure they are authenticated (which populates the token cache if needed)
        if (isAuthenticated(phoneNumber)) {
            return sessionTokens.get(phoneNumber);
        }
        return null;
    }

    public void logoutUser(String phoneNumber) {
        String token = sessionTokens.remove(phoneNumber);
        if (token != null) {
            imsApiService.logout(token);
        }
        Optional<UserEntity> userOpt = userRepository.findByPhoneNumber(phoneNumber);
        if (userOpt.isPresent()) {
            UserEntity user = userOpt.get();
            user.setEncryptedPasscode(null);
            userRepository.save(user);
        }
        studentNames.remove(phoneNumber);
        resetState(phoneNumber);
    }

    public String getStudentName(String phoneNumber, String token) {
        if (studentNames.containsKey(phoneNumber)) {
            return studentNames.get(phoneNumber);
        }
        
        // Try getting name from database first
        Optional<UserEntity> userOpt = userRepository.findByPhoneNumber(phoneNumber);
        if (userOpt.isPresent() && userOpt.get().getName() != null && !userOpt.get().getName().isEmpty()) {
            String dbName = userOpt.get().getName();
            studentNames.put(phoneNumber, dbName);
            return dbName;
        }

        // Fallback to IMS API
        String name;
        if (isFaculty(phoneNumber)) {
            name = imsApiService.getFacultyName(token);
        } else {
            name = imsApiService.getStudentName(token);
        }
        studentNames.put(phoneNumber, name);
        return name;
    }

    private final Map<String, List<Map<String, Object>>> attendanceSubjectParams = new ConcurrentHashMap<>();
    private final Map<String, Map<String, Object>> activeAttendanceParams = new ConcurrentHashMap<>();
    private final Map<String, List<com.rit.whatsapp.util.MessageFormatter.StudentAttendanceInfo>> activeStudentLists = new ConcurrentHashMap<>();

    public void storeAttendanceSubjectParams(String phoneNumber, List<Map<String, Object>> paramsList) {
        attendanceSubjectParams.put(phoneNumber, paramsList);
    }

    public Map<String, Object> getAttendanceSubjectParams(String phoneNumber, int index) {
        List<Map<String, Object>> paramsList = attendanceSubjectParams.get(phoneNumber);
        if (paramsList != null && index >= 0 && index < paramsList.size()) {
            return paramsList.get(index);
        }
        return null;
    }

    public void setActiveAttendanceParams(String phoneNumber, Map<String, Object> params) {
        activeAttendanceParams.put(phoneNumber, params);
    }

    public Map<String, Object> getActiveAttendanceParams(String phoneNumber) {
        return activeAttendanceParams.get(phoneNumber);
    }

    public void clearActiveAttendanceParams(String phoneNumber) {
        activeAttendanceParams.remove(phoneNumber);
        activeStudentLists.remove(phoneNumber);
    }

    public void storeActiveStudentList(String phoneNumber, List<com.rit.whatsapp.util.MessageFormatter.StudentAttendanceInfo> list) {
        activeStudentLists.put(phoneNumber, list);
    }

    public List<com.rit.whatsapp.util.MessageFormatter.StudentAttendanceInfo> getActiveStudentList(String phoneNumber) {
        return activeStudentLists.get(phoneNumber);
    }

    private final Map<String, Map<String, String>> receiptIdMappings = new ConcurrentHashMap<>();

    public void storeReceiptIdMappings(String phoneNumber, Map<String, String> mappings) {
        receiptIdMappings.put(phoneNumber, mappings);
    }

    public String getReceiptIdMapping(String phoneNumber, String shortId) {
        Map<String, String> mappings = receiptIdMappings.get(phoneNumber);
        if (mappings != null) {
            return mappings.get(shortId);
        }
        return null;
    }

    public void setTempLoginId(String phoneNumber, String tempId) {
        temporaryLoginIds.put(phoneNumber, tempId);
    }

    public String getTempLoginId(String phoneNumber) {
        return temporaryLoginIds.get(phoneNumber);
    }

    public void clearTempLoginId(String phoneNumber) {
        temporaryLoginIds.remove(phoneNumber);
    }

    private final Map<String, GrievanceDraft> grievanceDrafts = new ConcurrentHashMap<>();

    public GrievanceDraft getOrCreateGrievanceDraft(String phoneNumber) {
        return grievanceDrafts.computeIfAbsent(phoneNumber, k -> {
            GrievanceDraft draft = new GrievanceDraft();
            draft.setPhoneNumber(phoneNumber);
            return draft;
        });
    }

    public GrievanceDraft getGrievanceDraft(String phoneNumber) {
        return grievanceDrafts.get(phoneNumber);
    }

    public void clearGrievanceDraft(String phoneNumber) {
        grievanceDrafts.remove(phoneNumber);
    }
    
    public String getImsId(String phoneNumber) {
        Optional<UserEntity> userOpt = userRepository.findByPhoneNumber(phoneNumber);
        if (userOpt.isPresent()) {
            return userOpt.get().getImsId();
        }
        return null;
    }

    public boolean isFaculty(String phoneNumber) {
        Optional<UserEntity> userOpt = userRepository.findByPhoneNumber(phoneNumber);
        if (userOpt.isPresent()) {
            return userOpt.get().isFaculty();
        }
        return false;
    }

    public void deleteUser(String phoneNumber) {
        logoutUser(phoneNumber);
        Optional<UserEntity> userOpt = userRepository.findByPhoneNumber(phoneNumber);
        if (userOpt.isPresent()) {
            userRepository.delete(userOpt.get());
        }
    }
}