package com.writeoff.module.notification.ws;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Service
public class NotificationWebSocketPushService {
    private final JdbcTemplate jdbcTemplate;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final Map<String, CopyOnWriteArraySet<WebSocketSession>> userSessions = new ConcurrentHashMap<String, CopyOnWriteArraySet<WebSocketSession>>();
    private final Map<Long, CopyOnWriteArraySet<WebSocketSession>> tenantSessions = new ConcurrentHashMap<Long, CopyOnWriteArraySet<WebSocketSession>>();
    private final Map<String, String> sessionUserKey = new ConcurrentHashMap<String, String>();
    private final Map<String, Long> sessionTenantKey = new ConcurrentHashMap<String, Long>();

    public NotificationWebSocketPushService(JdbcTemplate jdbcTemplate) {
        this.jdbcTemplate = jdbcTemplate;
    }

    public void register(WebSocketSession session, Long tenantId, Long userId) {
        if (session == null || tenantId == null || userId == null) {
            return;
        }
        String userKey = userKey(tenantId, userId);
        userSessions.computeIfAbsent(userKey, k -> new CopyOnWriteArraySet<WebSocketSession>()).add(session);
        tenantSessions.computeIfAbsent(tenantId, k -> new CopyOnWriteArraySet<WebSocketSession>()).add(session);
        sessionUserKey.put(session.getId(), userKey);
        sessionTenantKey.put(session.getId(), tenantId);
    }

    public void unregister(WebSocketSession session) {
        if (session == null) {
            return;
        }
        String sid = session.getId();
        String userKey = sessionUserKey.remove(sid);
        Long tenantId = sessionTenantKey.remove(sid);
        if (userKey != null) {
            Set<WebSocketSession> set = userSessions.get(userKey);
            if (set != null) {
                set.remove(session);
                if (set.isEmpty()) {
                    userSessions.remove(userKey);
                }
            }
        }
        if (tenantId != null) {
            Set<WebSocketSession> set = tenantSessions.get(tenantId);
            if (set != null) {
                set.remove(session);
                if (set.isEmpty()) {
                    tenantSessions.remove(tenantId);
                }
            }
        }
    }

    public void pushInAppNotification(Long tenantId, String receiverRef, Long receiverUserId) {
        if (tenantId == null || tenantId <= 0) {
            return;
        }
        Long targetUserId = receiverUserId;
        if (targetUserId == null && receiverRef != null) {
            String ref = receiverRef.trim();
            if (ref.startsWith("user-")) {
                try {
                    targetUserId = Long.valueOf(ref.substring("user-".length()));
                } catch (Exception ignored) {
                }
            } else {
                try {
                    targetUserId = Long.valueOf(ref);
                } catch (Exception ignored) {
                }
            }
        }
        if ("ALL".equalsIgnoreCase(receiverRef == null ? "" : receiverRef.trim())) {
            pushToTenant(tenantId);
            return;
        }
        if (targetUserId != null && targetUserId > 0) {
            pushToUser(tenantId, targetUserId);
            return;
        }
        pushToTenant(tenantId);
    }

    private void pushToUser(Long tenantId, Long userId) {
        String key = userKey(tenantId, userId);
        Set<WebSocketSession> sessions = userSessions.get(key);
        if (sessions == null || sessions.isEmpty()) {
            return;
        }
        Map<String, Object> data = new LinkedHashMap<String, Object>();
        data.put("type", "IN_APP_UNREAD_CHANGED");
        data.put("tenantId", tenantId);
        data.put("userId", userId);
        data.put("unreadCount", countUnread(tenantId, userId));
        broadcast(sessions, data);
    }

    private void pushToTenant(Long tenantId) {
        Set<WebSocketSession> sessions = tenantSessions.get(tenantId);
        if (sessions == null || sessions.isEmpty()) {
            return;
        }
        Map<String, Object> data = new LinkedHashMap<String, Object>();
        data.put("type", "IN_APP_NOTIFICATION_NEW");
        data.put("tenantId", tenantId);
        broadcast(sessions, data);
    }

    private int countUnread(Long tenantId, Long userId) {
        Integer count = jdbcTemplate.queryForObject(
                "SELECT COUNT(1) FROM in_app_notification " +
                        "WHERE tenant_id=? AND is_deleted=0 AND status='UNREAD' AND (receiver_ref='ALL' OR receiver_ref=? OR receiver_user_id=?)",
                Integer.class,
                tenantId,
                "user-" + userId,
                userId
        );
        return count == null ? 0 : count;
    }

    private void broadcast(Set<WebSocketSession> sessions, Map<String, Object> payload) {
        String message;
        try {
            message = objectMapper.writeValueAsString(payload);
        } catch (Exception ex) {
            return;
        }
        for (WebSocketSession session : sessions) {
            if (session == null || !session.isOpen()) {
                continue;
            }
            try {
                session.sendMessage(new TextMessage(message));
            } catch (IOException ignored) {
            }
        }
    }

    private String userKey(Long tenantId, Long userId) {
        return String.valueOf(tenantId) + ":" + String.valueOf(userId);
    }
}
