package com.dji.sample.component.websocket.config;

import com.dji.sample.common.model.CustomClaim;
import com.dji.sample.common.util.JwtUtil;
import com.dji.sample.component.websocket.service.IWebSocketManageService;
import com.dji.sdk.websocket.WebSocketDefaultHandler;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

import java.security.Principal;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 *
 * @author sean.zhou
 * @date 2021/11/16
 * @version 0.1
 */
@Slf4j
public class MyWebSocketHandler extends WebSocketDefaultHandler {

    private IWebSocketManageService webSocketManageService;

    private ObjectMapper objectMapper = new ObjectMapper();
    private Map<String, Boolean> authenticatedSessions = new ConcurrentHashMap<>();
    private Map<String, CustomClaim> sessionClaims = new ConcurrentHashMap<>();

    MyWebSocketHandler(WebSocketHandler delegate, IWebSocketManageService webSocketManageService) {
        super(delegate);
        this.webSocketManageService = webSocketManageService;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Principal principal = session.getPrincipal();
        String principalName = principal.getName();

        if (StringUtils.hasText(principalName) && !principalName.startsWith("temp-")) {
            webSocketManageService.put(principalName, new MyConcurrentWebSocketSession(session));
            authenticatedSessions.put(session.getId(), true);
            log.debug("{} is connected (pre-authenticated). ID: {}. WebSocketSession[current count: {}]",
                    principalName, session.getId(), webSocketManageService.getConnectedCount());
        } else {
            log.debug("Unauthenticated connection established. ID: {}, temp principal: {}",
                    session.getId(), principalName);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String sessionId = session.getId();
        Boolean isAuthenticated = authenticatedSessions.get(sessionId);

        if (Boolean.TRUE.equals(isAuthenticated)) {
            CustomClaim claim = sessionClaims.get(sessionId);
            if (claim != null) {
                String key = claim.getWorkspaceId() + "/" + claim.getUserType() + "/" + claim.getId();
                webSocketManageService.remove(key, sessionId);
                log.debug("{} is disconnected. ID: {}. WebSocketSession[current count: {}]",
                        key, sessionId, webSocketManageService.getConnectedCount());
            }
        }

        authenticatedSessions.remove(sessionId);
        sessionClaims.remove(sessionId);

    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String sessionId = session.getId();
        Boolean isAuthenticated = authenticatedSessions.get(sessionId);

        String payload = message.getPayload().toString();
        log.debug("Received message from session {}: {}", sessionId, payload);

        try {
            JsonNode jsonNode = objectMapper.readTree(payload);
            String type = jsonNode.has("type") ? jsonNode.get("type").asText() : null;

            if ("auth".equals(type) && !Boolean.TRUE.equals(isAuthenticated)) {
                String token = jsonNode.has("token") ? jsonNode.get("token").asText() : null;
                if (StringUtils.hasText(token)) {
                    Optional<CustomClaim> customClaimOpt = JwtUtil.parseToken(token);
                    if (customClaimOpt.isPresent()) {
                        CustomClaim claim = customClaimOpt.get();
                        String key = claim.getWorkspaceId() + "/" + claim.getUserType() + "/" + claim.getId();

                        authenticatedSessions.put(sessionId, true);
                        sessionClaims.put(sessionId, claim);
                        webSocketManageService.put(key, new MyConcurrentWebSocketSession(session));

                        log.debug("Session {} authenticated successfully. User: {}", sessionId, key);

                        String authResponse = objectMapper.writeValueAsString(Map.of(
                                "type", "auth_success",
                                "status", "ok"
                        ));
                        session.sendMessage(new org.springframework.web.socket.TextMessage(authResponse));
                        return;
                    }
                }

                log.warn("Authentication failed for session {}", sessionId);
                String authFail = objectMapper.writeValueAsString(Map.of(
                        "type", "auth_fail",
                        "status", "error",
                        "message", "Invalid token"
                ));
                session.sendMessage(new org.springframework.web.socket.TextMessage(authFail));
                session.close(CloseStatus.NOT_ACCEPTABLE);
                return;
            }

            if (!Boolean.TRUE.equals(isAuthenticated)) {
                log.warn("Unauthenticated session {} tried to send message. Closing.", sessionId);
                String authRequired = objectMapper.writeValueAsString(Map.of(
                        "type", "auth_required",
                        "message", "Please authenticate first"
                ));
                session.sendMessage(new org.springframework.web.socket.TextMessage(authRequired));
//                session.close(CloseStatus.NOT_ACCEPTABLE);
                return;
            }

            super.handleMessage(session, message);

        } catch (Exception e) {
            log.error("Error handling message from session {}", sessionId, e);
            if (!Boolean.TRUE.equals(isAuthenticated)) {
                session.close(CloseStatus.BAD_DATA);
            }
        }
    }

}