package com.writeoff.module.notification.ws;

import com.writeoff.security.AuthScope;
import com.writeoff.security.JwtTokenService;
import io.jsonwebtoken.Claims;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.net.URI;
import java.util.LinkedHashMap;
import java.util.Map;

@Component
public class NotificationWebSocketHandler extends TextWebSocketHandler {
    private final JwtTokenService jwtTokenService;
    private final NotificationWebSocketPushService pushService;

    public NotificationWebSocketHandler(JwtTokenService jwtTokenService, NotificationWebSocketPushService pushService) {
        this.jwtTokenService = jwtTokenService;
        this.pushService = pushService;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String token = resolveToken(session);
        if (token == null || token.trim().isEmpty()) {
            session.close(CloseStatus.POLICY_VIOLATION);
            return;
        }
        try {
            Claims claims = jwtTokenService.parse(token.trim());
            AuthScope scope = AuthScope.fromClaim(claims.get("scope", String.class));
            if (scope != AuthScope.TENANT) {
                session.close(CloseStatus.POLICY_VIOLATION);
                return;
            }
            Number uid = claims.get("uid", Number.class);
            Number tid = claims.get("tid", Number.class);
            if (uid == null || tid == null) {
                session.close(CloseStatus.POLICY_VIOLATION);
                return;
            }
            Long userId = uid.longValue();
            Long tenantId = tid.longValue();
            session.getAttributes().put("uid", userId);
            session.getAttributes().put("tid", tenantId);
            pushService.register(session, tenantId, userId);
            Map<String, Object> hello = new LinkedHashMap<String, Object>();
            hello.put("type", "WS_CONNECTED");
            hello.put("tenantId", tenantId);
            hello.put("userId", userId);
            session.sendMessage(new TextMessage(new com.fasterxml.jackson.databind.ObjectMapper().writeValueAsString(hello)));
        } catch (Exception ex) {
            session.close(CloseStatus.POLICY_VIOLATION);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        // no-op: server push only
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        pushService.unregister(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        pushService.unregister(session);
    }

    private String resolveToken(WebSocketSession session) {
        URI uri = session.getUri();
        if (uri == null || uri.getQuery() == null) {
            return null;
        }
        String[] parts = uri.getQuery().split("&");
        for (String item : parts) {
            int idx = item.indexOf('=');
            if (idx <= 0) {
                continue;
            }
            String key = item.substring(0, idx);
            if (!"token".equals(key)) {
                continue;
            }
            String val = item.substring(idx + 1);
            try {
                return java.net.URLDecoder.decode(val, "UTF-8");
            } catch (Exception ex) {
                return val;
            }
        }
        return null;
    }
}
