package com.gnx.telemedicine.security;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

import java.security.Principal;

@Configuration
@EnableWebSocketMessageBroker
@Order(Ordered.HIGHEST_PRECEDENCE + 99)
@RequiredArgsConstructor
@Slf4j
public class WebSocketSecurityConfig implements WebSocketMessageBrokerConfigurer {

    private final JwtUtils jwtUtils;
    private final CustomUserDetailsService customUserDetailsService;

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(new ChannelInterceptor() {
            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {
                StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
                
                if (accessor != null) {
                    StompCommand command = accessor.getCommand();
                    String destination = accessor.getDestination();
                    
                    // Log ALL SEND commands with full details
                    if (StompCommand.SEND.equals(command)) {
                        log.debug("========================================");
                        log.debug(">>> SEND COMMAND INTERCEPTED");
                        log.debug(">>> Destination: {}", destination);
                        log.debug(">>> Principal BEFORE auth: {}", accessor.getUser() != null ? accessor.getUser().getName() : "NULL");
                        log.debug(">>> Command: {}", command);
                        log.debug(">>> Session ID: {}", accessor.getSessionId());
                        
                        // Log message payload
                        try {
                            Object payload = message.getPayload();
                            log.debug(">>> Payload type: {}", payload != null ? payload.getClass().getName() : "NULL");
                            if (payload instanceof byte[]) {
                                String bodyStr = new String((byte[]) payload);
                                log.debug(">>> Payload body: {}", bodyStr);
                            } else {
                                log.debug(">>> Payload: {}", payload);
                            }
                        } catch (Exception e) {
                            log.warn(">>> Error reading payload: {}", e.getMessage());
                        }
                        
                        // Log all native headers
                        java.util.Map<String, java.util.List<String>> nativeHeaders = accessor.toNativeHeaderMap();
                        if (nativeHeaders != null && !nativeHeaders.isEmpty()) {
                            log.debug(">>> Native headers: {}", nativeHeaders);
                        } else {
                            log.debug(">>> No native headers found");
                        }
                        
                        if (destination != null && destination.contains("call")) {
                            log.debug(">>> *** CALL MESSAGE TYPE ***");
                            if (destination.contains("initiate")) {
                                log.debug(">>> *** CALL.INITIATE SPECIFICALLY ***");
                            }
                        }
                        log.debug("========================================");
                    }
                    
                    // Also log CONNECT commands to verify interceptor is working
                    if (StompCommand.CONNECT.equals(command)) {
                        log.debug(">>> CONNECT command detected");
                    }
                    
                    // Authenticate on CONNECT
                    if (StompCommand.CONNECT.equals(command)) {
                        String authToken = accessor.getFirstNativeHeader("Authorization");
                        log.debug("CONNECT command - Auth token: {}", authToken != null ? "PRESENT" : "NULL");
                        
                        if (authToken != null && authToken.startsWith("Bearer ")) {
                            String token = authToken.substring(7);
                            
                            if (jwtUtils.validateToken(token)) {
                                String email = jwtUtils.getEmailFromToken(token);
                                UserDetails userDetails = customUserDetailsService.loadUserByUsername(email);
                                
                                if (userDetails.isEnabled()) {
                                    UsernamePasswordAuthenticationToken auth = new UsernamePasswordAuthenticationToken(
                                            email,
                                            null,
                                            userDetails.getAuthorities()
                                    );
                                    accessor.setUser(auth);
                                    
                                    // Store token and Principal in session for later use
                                    java.util.Map<String, Object> sessionAttrs = accessor.getSessionAttributes();
                                    if (sessionAttrs != null) {
                                        sessionAttrs.put("authToken", authToken);
                                        sessionAttrs.put("principal", auth);
                                        log.debug("Stored Principal in session attributes");
                                    } else {
                                        log.warn("WARNING: Session attributes are null during CONNECT");
                                    }
                                    
                                    log.info("CONNECT authenticated user: {}", email);
                                } else {
                                    log.warn("ERROR: User account is disabled: {}", email);
                                }
                            } else {
                                log.warn("ERROR: Invalid JWT token on CONNECT");
                            }
                        } else {
                            log.warn("ERROR: No auth token found on CONNECT");
                        }
                    }
                    // For SEND commands, verify Principal exists (should be set from CONNECT)
                    else if (StompCommand.SEND.equals(command)) {
                        Principal principal = accessor.getUser();
                        if (destination != null && destination.contains("call")) {
                            log.debug("=== SEND MESSAGE TO: {} ===", destination);
                            log.debug("Principal: {}", principal != null ? principal.getName() : "NULL");
                        }
                        if (principal == null) {
                            log.warn("WARNING: Principal is NULL for SEND command to: {}", destination);
                            
                            // First, try to get Principal from session attributes (set during CONNECT)
                            java.util.Map<String, Object> sessionAttrs = accessor.getSessionAttributes();
                            if (sessionAttrs != null) {
                                Object storedPrincipal = sessionAttrs.get("principal");
                                if (storedPrincipal instanceof UsernamePasswordAuthenticationToken) {
                                    accessor.setUser((Principal) storedPrincipal);
                                    log.debug("Principal restored from session attributes: {}", ((Principal) storedPrincipal).getName());
                                    principal = (Principal) storedPrincipal;
                                }
                            }
                            
                            // If still null, try to authenticate from token
                            if (principal == null) {
                                String authToken = accessor.getFirstNativeHeader("Authorization");
                                log.debug("Auth token from headers: {}", authToken != null ? "PRESENT" : "NULL");
                                
                                // Also try to get token from session attributes (set during CONNECT)
                                if (authToken == null && sessionAttrs != null) {
                                    Object sessionToken = sessionAttrs.get("authToken");
                                    if (sessionToken instanceof String) {
                                        authToken = (String) sessionToken;
                                        log.debug("Found auth token in session attributes");
                                    }
                                }
                                
                                if (authToken != null && authToken.startsWith("Bearer ")) {
                                    String token = authToken.substring(7);
                                    if (jwtUtils.validateToken(token)) {
                                        String email = jwtUtils.getEmailFromToken(token);
                                        UserDetails userDetails = customUserDetailsService.loadUserByUsername(email);
                                        if (userDetails.isEnabled()) {
                                            UsernamePasswordAuthenticationToken auth = new UsernamePasswordAuthenticationToken(
                                                    email,
                                                    null,
                                                    userDetails.getAuthorities()
                                            );
                                            accessor.setUser(auth);
                                            // Store in session for future messages
                                            if (sessionAttrs != null) {
                                                sessionAttrs.put("principal", auth);
                                            }
                                            log.debug("Principal set from token: {}", email);
                                            principal = auth;
                                        } else {
                                            log.warn("ERROR: User account is disabled: {}", email);
                                        }
                                    } else {
                                        log.warn("ERROR: Invalid JWT token");
                                    }
                                } else {
                                    log.warn("ERROR: No valid auth token found for SEND command");
                                }
                            }
                        } else {
                            log.debug("Principal is set: {}", principal.getName());
                        }
                        
                        // Log Principal AFTER authentication attempt
                        Principal finalPrincipal = accessor.getUser();
                        if (destination != null && destination.contains("call")) {
                            log.debug(">>> FINAL Principal for call message: {}", finalPrincipal != null ? finalPrincipal.getName() : "NULL");
                        }
                    }
                }
                return message;
            }
        });
    }
}


