package com.petshop.backend.config;

import com.petshop.backend.entity.User;
import com.petshop.backend.repository.UserRepository;
import com.petshop.backend.security.AppPrincipal;
import com.petshop.backend.security.JwtUtil;
import com.petshop.backend.service.ChatService;
import io.jsonwebtoken.JwtException;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.stereotype.Component;

import java.security.Principal;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;

@Component
public class WebSocketAuthChannelInterceptor implements ChannelInterceptor {

    private final JwtUtil jwtUtil;
    private final UserRepository userRepository;
    private final ChatService chatService;

    public WebSocketAuthChannelInterceptor(JwtUtil jwtUtil, UserRepository userRepository, ChatService chatService) {
        this.jwtUtil = jwtUtil;
        this.userRepository = userRepository;
        this.chatService = chatService;
    }

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        StompCommand command = accessor.getCommand();

        if (command == null) {
            return message;
        }

        if (StompCommand.CONNECT.equals(command)) {
            String tokenHeader = firstHeader(accessor, "Authorization");
            String token = extractToken(tokenHeader != null ? tokenHeader : firstHeader(accessor, "token"));
            if (token == null || token.isBlank()) {
                throw new IllegalArgumentException("Missing websocket token");
            }

            Long userId = extractUserId(token);
            User user = userId == null ? null : userRepository.findById(userId).orElse(null);
            if (user == null) {
                throw new IllegalArgumentException("User not found");
            }
            if (user.getActive() == null || !user.getActive()) {
                throw new IllegalArgumentException("User account is inactive");
            }
            if (!jwtUtil.validateToken(token, user)) {
                throw new IllegalArgumentException("Invalid websocket token");
            }

            AppPrincipal principal = new AppPrincipal(
                    user.getId(),
                    user.getUsername(),
                    user.getRole(),
                    user.getTokenVersion()
            );
            UsernamePasswordAuthenticationToken authentication = new UsernamePasswordAuthenticationToken(
                    principal,
                    null,
                    principal.getAuthorities()
            );
            accessor.setUser(authentication);
            accessor.getSessionAttributes().put("user", authentication);
            return message;
        }

        if (StompCommand.DISCONNECT.equals(command) || StompCommand.UNSUBSCRIBE.equals(command)) {
            return message;
        }

        User user = resolveUser(accessor.getUser(), accessor);
        if (user == null) {
            throw new IllegalArgumentException("Unauthenticated websocket session");
        }

        if (StompCommand.SUBSCRIBE.equals(command)) {
            authorizeSubscription(accessor.getDestination(), user);
        } else if (StompCommand.SEND.equals(command)) {
            authorizeSend(accessor.getDestination(), user);
        }

        return message;
    }

    private User resolveUser(Principal principal, StompHeaderAccessor accessor) {
        Principal currentPrincipal = principal;
        if (currentPrincipal == null && accessor.getSessionAttributes() != null) {
            Object sessionUser = accessor.getSessionAttributes().get("user");
            if (sessionUser instanceof Principal storedPrincipal) {
                accessor.setUser(storedPrincipal);
                currentPrincipal = storedPrincipal;
            }
        }

        if (currentPrincipal instanceof UsernamePasswordAuthenticationToken authenticationToken
                && authenticationToken.getPrincipal() instanceof AppPrincipal appPrincipal) {
            return userRepository.findById(appPrincipal.getUserId())
                    .orElseThrow(() -> new IllegalArgumentException("User not found"));
        }

        if (currentPrincipal instanceof AppPrincipal appPrincipal) {
            return userRepository.findById(appPrincipal.getUserId())
                    .orElseThrow(() -> new IllegalArgumentException("User not found"));
        }

        String tokenHeader = firstHeader(accessor, "Authorization");
        String token = extractToken(tokenHeader != null ? tokenHeader : firstHeader(accessor, "token"));
        if (token == null || token.isBlank()) {
            return null;
        }

        Long userId = extractUserId(token);
        User user = userId == null ? null : userRepository.findById(userId).orElse(null);
        if (user == null) {
            throw new IllegalArgumentException("User not found");
        }
        if (user.getActive() == null || !user.getActive()) {
            throw new IllegalArgumentException("User account is inactive");
        }
        if (!jwtUtil.validateToken(token, user)) {
            throw new IllegalArgumentException("Invalid websocket token");
        }
        return user;
    }

    private void authorizeSubscription(String destination, User user) {
        destination = normalizeDestination(destination);
        if (destination == null || destination.startsWith("/user/queue/")) {
            return;
        }

        if ("/topic/chat/conversations".equals(destination)) {
            if (user.getRole() == User.Role.CUSTOMER) {
                throw new IllegalArgumentException("Customers cannot subscribe to staff conversation feed");
            }
            return;
        }

        Long conversationId = extractConversationId(destination, "/topic/chat/conversations/");
        if (conversationId != null && chatService.hasConversationAccess(conversationId, user.getId(), user.getRole())) {
            return;
        }

        throw new IllegalArgumentException("Not authorized to subscribe to destination");
    }

    private void authorizeSend(String destination, User user) {
        destination = normalizeDestination(destination);
        Long conversationId = extractConversationId(destination, "/app/chat/conversations/");
        if (conversationId != null && destination.endsWith("/messages") && chatService.hasConversationAccess(conversationId, user.getId(), user.getRole())) {
            return;
        }

        throw new IllegalArgumentException("Not authorized to send to destination");
    }

    private Long extractConversationId(String destination, String prefix) {
        if (destination == null || !destination.startsWith(prefix)) {
            return null;
        }

        String suffix = destination.substring(prefix.length());
        String[] parts = suffix.split("/");
        if (parts.length == 0 || parts[0].isBlank()) {
            return null;
        }

        try {
            return Long.parseLong(parts[0]);
        } catch (NumberFormatException ex) {
            return null;
        }
    }

    private String firstHeader(StompHeaderAccessor accessor, String name) {
        List<String> values = accessor.getNativeHeader(name);
        if (values != null && !values.isEmpty()) {
            return values.get(0);
        }
        for (String headerName : accessor.toNativeHeaderMap().keySet()) {
            if (headerName.equalsIgnoreCase(name)) {
                List<String> alternateValues = accessor.getNativeHeader(headerName);
                return alternateValues == null || alternateValues.isEmpty() ? null : alternateValues.get(0);
            }
        }
        return null;
    }

    private String extractToken(String rawValue) {
        if (rawValue == null || rawValue.isBlank()) {
            return null;
        }
        String normalized = rawValue.trim();
        return normalized.regionMatches(true, 0, "Bearer ", 0, 7) ? normalized.substring(7) : normalized;
    }

    private String normalizeDestination(String destination) {
        if (destination == null || destination.isBlank()) {
            return destination;
        }
        String normalized = destination.replace('\\', '/');
        while (normalized.contains("//")) {
            normalized = normalized.replace("//", "/");
        }
        return normalized.toLowerCase(Locale.ROOT);
    }

    private Long extractUserId(String token) {
        try {
            return jwtUtil.extractUserId(token);
        } catch (JwtException | IllegalArgumentException ex) {
            throw new IllegalArgumentException("Invalid websocket token: " + ex.getMessage(), ex);
        }
    }

    public Map<String, Object> buildErrorPayload(Exception ex, String destination, Principal principal) {
        Map<String, Object> response = new LinkedHashMap<>();
        response.put("message", ex.getMessage() == null || ex.getMessage().isBlank() ? "WebSocket request failed" : ex.getMessage());
        response.put("details", ex.getClass().getSimpleName());
        response.put("destination", normalizeDestination(destination));
        response.put("authenticated", principal != null);
        return response;
    }
}
