package de.assecutor.votianlt.messaging;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
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.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.*;

/**
 * WebSocket service for direct bidirectional communication with mobile clients.
 *
 * Wire Protocol: Each WebSocket message is a JSON document with a "topic" and
 * "payload" field:
 *
 * <pre>
 * {
 *   "topic": "/server/login",
 *   "payload": { ... }
 * }
 * </pre>
 *
 * Topic Structure:
 * <ul>
 * <li>Server to Client: /client/{messageType}</li>
 * <li>Client to Server: /server/{messageType}</li>
 * <li>Login (special): /server/login (unauthenticated)</li>
 * </ul>
 */
@Component
@Slf4j
public class WebSocketService extends TextWebSocketHandler {

    @FunctionalInterface
    public interface MessageHandler {
        void onMessageReceived(String clientId, byte[] payload);
    }

    private static final String TOPIC_TO_CLIENT = "/client/%s";
    private static final long PENDING_SESSION_TIMEOUT_MS = 30_000;

    private final ObjectMapper objectMapper;

    // appUserId -> WebSocketSession
    private final ConcurrentHashMap<String, WebSocketSession> clientSessions = new ConcurrentHashMap<>();

    // sessionId -> appUserId (reverse lookup for cleanup on disconnect)
    private final ConcurrentHashMap<String, String> sessionToClient = new ConcurrentHashMap<>();

    // sessionId -> PendingSession (connected but not yet logged in)
    private final ConcurrentHashMap<String, PendingSession> pendingSessions = new ConcurrentHashMap<>();

    private final Map<String, MessageHandler> messageHandlers = new ConcurrentHashMap<>();
    private volatile boolean initialized = false;

    private ScheduledExecutorService pendingSessionCleanup;

    public WebSocketService(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    // ==========================================
    // Lifecycle
    // ==========================================

    @PostConstruct
    public void init() {
        pendingSessionCleanup = Executors.newSingleThreadScheduledExecutor(r -> {
            Thread t = new Thread(r, "ws-pending-cleanup");
            t.setDaemon(true);
            return t;
        });
        pendingSessionCleanup.scheduleAtFixedRate(this::cleanupPendingSessions, 30, 30, TimeUnit.SECONDS);

        initialized = true;
        log.info("[WebSocket] Service initialized on endpoint /ws/messaging");
    }

    @PreDestroy
    public void shutdown() {
        if (pendingSessionCleanup != null) {
            pendingSessionCleanup.shutdownNow();
        }

        for (var entry : clientSessions.entrySet()) {
            try {
                WebSocketSession session = entry.getValue();
                if (session.isOpen()) {
                    session.close(CloseStatus.GOING_AWAY);
                }
            } catch (Exception e) {
                log.warn("[WebSocket] Error closing session for client {}: {}", entry.getKey(), e.getMessage());
            }
        }

        for (var entry : pendingSessions.entrySet()) {
            try {
                if (entry.getValue().session.isOpen()) {
                    entry.getValue().session.close(CloseStatus.GOING_AWAY);
                }
            } catch (Exception ignored) {
            }
        }

        clientSessions.clear();
        sessionToClient.clear();
        pendingSessions.clear();
        messageHandlers.clear();
        initialized = false;

        log.info("[WebSocket] Service shut down");
    }

    // ==========================================
    // Public API
    // ==========================================

    public CompletableFuture<Void> sendToClient(String clientId, String messageType, byte[] payload) {
        WebSocketSession session = clientSessions.get(clientId);
        if (session == null) {
            log.warn("[WebSocket] No session found for client {}", clientId);
            return CompletableFuture.failedFuture(new IOException("No WebSocket session for client: " + clientId));
        }

        if (!session.isOpen()) {
            log.warn("[WebSocket] Session for client {} is closed", clientId);
            // Session aus der Map entfernen
            clientSessions.remove(clientId);
            sessionToClient.remove(session.getId());
            return CompletableFuture.failedFuture(new IOException("WebSocket session closed for client: " + clientId));
        }

        try {
            String topic = String.format(TOPIC_TO_CLIENT, messageType);
            String payloadJson = new String(payload, StandardCharsets.UTF_8);

            ObjectNode wireMessage = objectMapper.createObjectNode();
            wireMessage.put("topic", topic);
            wireMessage.set("payload", objectMapper.readTree(payloadJson));

            String wireJson = objectMapper.writeValueAsString(wireMessage);
            log.info("[WebSocket OUT] {} to client {} (session open: {})", topic, clientId, session.isOpen());
            log.debug("[WebSocket OUT] {} -> {}", topic, wireJson);

            sendToSession(session, wireJson);
            log.debug("[WebSocket] Message sent successfully to client {}", clientId);
            return CompletableFuture.completedFuture(null);
        } catch (Exception e) {
            log.error("[WebSocket] Failed to send to client {}: {}", clientId, e.getMessage(), e);
            return CompletableFuture.failedFuture(new IOException("Failed to send WebSocket message", e));
        }
    }

    public void registerMessageHandler(String messageType, MessageHandler handler) {
        messageHandlers.put(messageType, handler);
        log.debug("[WebSocket] Registered handler for messageType: {}", messageType);
    }

    public boolean isConnected() {
        return initialized;
    }

    public boolean isClientConnected(String clientId) {
        WebSocketSession session = clientSessions.get(clientId);
        return session != null && session.isOpen();
    }

    public int getConnectedClientCount() {
        return clientSessions.size();
    }

    // ==========================================
    // WebSocket handler methods
    // ==========================================

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        pendingSessions.put(session.getId(), new PendingSession(session, Instant.now()));
        log.info("[WebSocket] New connection: sessionId={}, remote={}", session.getId(), session.getRemoteAddress());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        try {
            String json = message.getPayload();
            JsonNode wireMessage = objectMapper.readTree(json);

            JsonNode topicNode = wireMessage.get("topic");
            JsonNode payloadNode = wireMessage.get("payload");

            if (topicNode == null || payloadNode == null) {
                log.warn("[WebSocket] Invalid message format (missing topic or payload): {}", json);
                return;
            }

            String topic = topicNode.asText();
            byte[] payloadBytes = objectMapper.writeValueAsBytes(payloadNode);

            log.info("[WebSocket IN] {} <- {}", topic, json);

            // Login message (special: unauthenticated)
            if ("/server/login".equals(topic)) {
                handleLoginMessage(session, payloadBytes);
                return;
            }

            // Regular client message: /server/{messageType}
            if (topic.startsWith("/server/")) {
                // Verify session is authenticated
                String appUserId = sessionToClient.get(session.getId());
                if (appUserId == null) {
                    log.warn("[WebSocket] Unauthenticated session {} tried to send: {}", session.getId(), topic);
                    return;
                }
                handleClientMessage(topic, appUserId, payloadBytes);
            }
        } catch (Exception e) {
            log.error("[WebSocket] Error handling message: {}", e.getMessage(), e);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String sessionId = session.getId();

        // Remove from pending sessions
        pendingSessions.remove(sessionId);

        // Remove from authenticated sessions
        String clientId = sessionToClient.remove(sessionId);
        if (clientId != null) {
            clientSessions.remove(clientId, session);
            log.info("[WebSocket] Client disconnected: clientId={}, reason={}", clientId, status);
        } else {
            log.info("[WebSocket] Unauthenticated session closed: sessionId={}", sessionId);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        log.error("[WebSocket] Transport error for session {}: {}", session.getId(), exception.getMessage());
    }

    // ==========================================
    // Internal message routing
    // ==========================================

    private void handleLoginMessage(WebSocketSession session, byte[] payloadBytes) {
        MessageHandler handler = messageHandlers.get("login");
        if (handler != null) {
            handler.onMessageReceived(session.getId(), payloadBytes);
        }
    }

    /**
     * Register a pending session as authenticated under the given appUserId. Called
     * by MessagingConfig after successful login.
     */
    public void registerAuthenticatedSession(String wsSessionId, String appUserId) {
        PendingSession pending = pendingSessions.get(wsSessionId);
        if (pending == null) {
            log.warn("[WebSocket] No pending session for wsSessionId={}", wsSessionId);
            return;
        }
        registerClientSession(appUserId, pending.session());
    }

    /**
     * Send a wire-format message directly to a session by its WebSocket sessionId.
     * Used for sending login responses to pending (not yet authenticated) sessions.
     */
    public void sendToSessionById(String wsSessionId, String topic, byte[] payload) {
        try {
            // Check pending sessions first
            PendingSession pending = pendingSessions.get(wsSessionId);
            WebSocketSession session = pending != null ? pending.session() : null;

            // Fallback: check authenticated sessions via reverse lookup
            if (session == null) {
                String appUserId = sessionToClient.get(wsSessionId);
                if (appUserId != null) {
                    session = clientSessions.get(appUserId);
                }
            }

            if (session == null || !session.isOpen()) {
                log.warn("[WebSocket] Cannot send to session {}: not found or closed", wsSessionId);
                return;
            }

            String payloadJson = new String(payload, StandardCharsets.UTF_8);
            ObjectNode wireMessage = objectMapper.createObjectNode();
            wireMessage.put("topic", topic);
            wireMessage.set("payload", objectMapper.readTree(payloadJson));

            String wireJson = objectMapper.writeValueAsString(wireMessage);
            log.info("[WebSocket OUT] {} -> {}", topic, wireJson);

            sendToSession(session, wireJson);
        } catch (Exception e) {
            log.error("[WebSocket] Error sending to session {}: {}", wsSessionId, e.getMessage());
        }
    }

    private void handleClientMessage(String topic, String appUserId, byte[] payload) {
        String[] parts = topic.split("/");

        // Handle /server/{messageType} where messageType can contain slashes
        if (parts.length >= 3) {
            String messageType = String.join("/", Arrays.copyOfRange(parts, 2, parts.length));

            MessageHandler handler = messageHandlers.get(messageType);
            if (handler != null) {
                handler.onMessageReceived(appUserId, payload);
            } else {
                log.warn("[WebSocket] No handler registered for messageType: {}", messageType);
            }
        }
    }

    // ==========================================
    // Session management
    // ==========================================

    private void registerClientSession(String clientId, WebSocketSession session) {
        // Close old session if same clientId reconnects
        WebSocketSession oldSession = clientSessions.put(clientId, session);
        if (oldSession != null && oldSession.isOpen() && !oldSession.getId().equals(session.getId())) {
            try {
                String oldSessionId = oldSession.getId();
                sessionToClient.remove(oldSessionId);
                oldSession.close(CloseStatus.NORMAL.withReason("Replaced by new connection"));
                log.info("[WebSocket] Closed old session for clientId={} (replaced)", clientId);
            } catch (IOException e) {
                log.warn("[WebSocket] Error closing old session for client {}: {}", clientId, e.getMessage());
            }
        }

        sessionToClient.put(session.getId(), clientId);
        pendingSessions.remove(session.getId());

        log.info("[WebSocket] Client registered: clientId={}, sessionId={}", clientId, session.getId());
    }

    private void cleanupPendingSessions() {
        Instant cutoff = Instant.now().minusMillis(PENDING_SESSION_TIMEOUT_MS);
        pendingSessions.entrySet().removeIf(entry -> {
            if (entry.getValue().connectedAt.isBefore(cutoff)) {
                try {
                    WebSocketSession session = entry.getValue().session;
                    if (session.isOpen()) {
                        session.close(CloseStatus.POLICY_VIOLATION.withReason("Login timeout"));
                    }
                    log.info("[WebSocket] Closed pending session (login timeout): sessionId={}", entry.getKey());
                } catch (IOException e) {
                    log.warn("[WebSocket] Error closing pending session: {}", e.getMessage());
                }
                return true;
            }
            return false;
        });
    }

    // ==========================================
    // Utility methods
    // ==========================================

    private void sendToSession(WebSocketSession session, String message) throws IOException {
        synchronized (session) {
            if (session.isOpen()) {
                session.sendMessage(new TextMessage(message));
            }
        }
    }

    // ==========================================
    // Internal types
    // ==========================================

    private record PendingSession(WebSocketSession session, Instant connectedAt) {
    }
}
