package com.example.petstoremobile.websocket;

import android.os.Handler;
import android.os.Looper;
import android.util.Log;
import com.example.petstoremobile.dtos.ConversationDTO;
import com.example.petstoremobile.dtos.MessageDTO;
import com.google.gson.Gson;
import java.util.HashMap;
import java.util.Collections;
import java.util.Locale;
import java.util.Map;
import io.reactivex.android.schedulers.AndroidSchedulers;
import io.reactivex.disposables.CompositeDisposable;
import io.reactivex.disposables.Disposable;
import io.reactivex.schedulers.Schedulers;
import ua.naiksoftware.stomp.Stomp;
import ua.naiksoftware.stomp.StompClient;
import ua.naiksoftware.stomp.dto.StompHeader;

//Used to handle the websocket connection for the chat
public class StompChatManager {

    private static final String TAG = "StompChatManager";

    //Interface for when a message is received
    public interface MessageListener {
        void onMessageReceived(MessageDTO message);
    }

    //Interface for when a conversation is created or updated
    public interface ConversationListener {
        void onConversationUpdated(ConversationDTO conversation);
    }

    //Interface for when the websocket connection is opened, closed, or has an error
    public interface ConnectionListener {
        void onSocketOpened();
        void onSocketClosed();
        void onSocketError();
    }

    private StompClient stompClient;
    private final CompositeDisposable compositeDisposable = new CompositeDisposable();
    private Disposable topicDisposable;
    private Disposable conversationsDisposable;
    private Disposable userConversationsDisposable;
    private Disposable errorQueueDisposable;
    private final Gson gson = new Gson();
    private final Handler reconnectHandler = new Handler(Looper.getMainLooper());
    private MessageListener messageListener;
    private ConversationListener conversationListener;
    private ConnectionListener connectionListener;
    private final String authToken;
    private final String role;
    private final String baseUrl;
    private boolean isConnected;
    private boolean isConnecting;
    private boolean manualDisconnect;
    private Long pendingConversationId;

    public StompChatManager(String authToken, String role, String baseUrl) {
        this.authToken = authToken;
        this.role = role == null ? "" : role.trim().toUpperCase(Locale.ROOT);
        this.baseUrl = baseUrl;
    }

    public void setMessageListener(MessageListener listener) {
        this.messageListener = listener;
    }

    public void setConversationListener(ConversationListener listener) {
        this.conversationListener = listener;
    }

    public void setConnectionListener(ConnectionListener listener) {
        this.connectionListener = listener;
    }

    // Set up a stomp connection
    public void connect() {
        if (authToken == null || authToken.isBlank()) {
            Log.e(TAG, "Cannot connect websocket without token");
            return;
        }
        if (isConnected || isConnecting) {
            return;
        }
        manualDisconnect = false;
        isConnecting = true;
        reconnectHandler.removeCallbacksAndMessages(null);
        String webSocketUrl = buildWebSocketUrl();

        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", "Bearer " + authToken);

        stompClient = Stomp.over(Stomp.ConnectionProvider.OKHTTP, webSocketUrl, headers);
        stompClient.withClientHeartbeat(0).withServerHeartbeat(0);

        compositeDisposable.add(
                stompClient.lifecycle()
                        .subscribeOn(Schedulers.io())
                        .observeOn(AndroidSchedulers.mainThread())
                        .subscribe(event -> {
                            switch (event.getType()) {
                                case OPENED:
                                    isConnected = true;
                                    isConnecting = false;
                                    Log.d(TAG, "WebSocket opened");
                                    if (connectionListener != null) {
                                        connectionListener.onSocketOpened();
                                    }
                                    subscribeToErrorQueue();
                                    subscribeToConversationFeeds();
                                    if (pendingConversationId != null) {
                                        subscribeToTopic(pendingConversationId);
                                    }
                                    break;
                                case CLOSED:
                                    isConnected = false;
                                    isConnecting = false;
                                    Log.d(TAG, "WebSocket closed");
                                    if (connectionListener != null) {
                                        connectionListener.onSocketClosed();
                                    }
                                    scheduleReconnect();
                                    break;
                                case ERROR:
                                    isConnected = false;
                                    isConnecting = false;
                                    Log.e(TAG, "WebSocket error: " + event.getException());
                                    if (connectionListener != null) {
                                        connectionListener.onSocketError();
                                    }
                                    scheduleReconnect();
                                    break;
                            }
                        })
        );

        stompClient.connect(Collections.singletonList(
                new StompHeader("Authorization", "Bearer " + authToken)
        ));
    }

    // Subscribes to updates for a specific conversation
    public void subscribeToConversation(Long conversationId) {
        pendingConversationId = conversationId;
        if (!isConnected || stompClient == null) {
            Log.d(TAG, "Delaying subscription until socket opens for conversation " + conversationId);
            connect();
            return;
        }
        subscribeToTopic(conversationId);
    }

    // Clears the current conversation subscription
    public void clearConversationSubscription() {
        pendingConversationId = null;
        if (topicDisposable != null && !topicDisposable.isDisposed()) {
            topicDisposable.dispose();
            topicDisposable = null;
        }
    }

    //helper function to subscribe to a specific conversation topic
    private void subscribeToTopic(Long conversationId) {
        if (topicDisposable != null && !topicDisposable.isDisposed()) {
            topicDisposable.dispose();
        }

        String topic = "/topic/chat/conversations/" + conversationId;
        Log.d(TAG, "Subscribing to topic " + topic);

        topicDisposable = stompClient.topic(topic)
                .subscribeOn(Schedulers.io())
                .observeOn(AndroidSchedulers.mainThread())
                .subscribe(
                        stompMessage -> {
                            MessageDTO message = gson.fromJson(
                                    stompMessage.getPayload(), MessageDTO.class);
                            if (messageListener != null) {
                                messageListener.onMessageReceived(message);
                            }
                        },
                        throwable -> Log.e(TAG, "Topic error", throwable)
                );

        compositeDisposable.add(topicDisposable);
    }

    // Listens for conversation updates and refresh the chat list
    private void subscribeToConversationFeeds() {
        if (conversationsDisposable != null && !conversationsDisposable.isDisposed()) {
            conversationsDisposable.dispose();
        }
        if (userConversationsDisposable != null && !userConversationsDisposable.isDisposed()) {
            userConversationsDisposable.dispose();
        }

        userConversationsDisposable = stompClient.topic("/user/queue/chat/conversations")
                .subscribeOn(Schedulers.io())
                .observeOn(AndroidSchedulers.mainThread())
                .subscribe(
                        stompMessage -> {
                            Log.d(TAG, "Queue conversation update: " + stompMessage.getPayload());
                            ConversationDTO conversation = gson.fromJson(
                                    stompMessage.getPayload(), ConversationDTO.class);
                            if (conversationListener != null) {
                                conversationListener.onConversationUpdated(conversation);
                            }
                        },
                        throwable -> Log.e(TAG, "Conversation queue error", throwable)
                );

        compositeDisposable.add(userConversationsDisposable);

        if (isCustomer()) {
            return;
        }

        conversationsDisposable = stompClient.topic("/topic/chat/conversations")
                .subscribeOn(Schedulers.io())
                .observeOn(AndroidSchedulers.mainThread())
                .subscribe(
                        stompMessage -> {
                            Log.d(TAG, "Broadcast conversation update: " + stompMessage.getPayload());
                            ConversationDTO conversation = gson.fromJson(
                                    stompMessage.getPayload(), ConversationDTO.class);
                            if (conversationListener != null) {
                                conversationListener.onConversationUpdated(conversation);
                            }
                        },
                        throwable -> Log.e(TAG, "Conversation topic error", throwable)
                );

        compositeDisposable.add(conversationsDisposable);
    }

    // Log any error from stomp
    private void subscribeToErrorQueue() {
        if (errorQueueDisposable != null && !errorQueueDisposable.isDisposed()) {
            errorQueueDisposable.dispose();
        }

        errorQueueDisposable = stompClient.topic("/user/queue/chat/errors")
                .subscribeOn(Schedulers.io())
                .observeOn(AndroidSchedulers.mainThread())
                .subscribe(
                        stompMessage -> Log.e(TAG, "WebSocket queue error payload: " + stompMessage.getPayload()),
                        throwable -> Log.e(TAG, "WebSocket error queue subscription failed", throwable)
                );

        compositeDisposable.add(errorQueueDisposable);
    }

    // Disconnects the stomp connection
    public void disconnect() {
        manualDisconnect = true;
        isConnected = false;
        isConnecting = false;
        pendingConversationId = null;
        reconnectHandler.removeCallbacksAndMessages(null);
        compositeDisposable.clear();
        if (stompClient != null) {
            stompClient.disconnect();
        }
    }

    // Make the URL for the websocket connection
    private String buildWebSocketUrl() {
        String cleanBaseUrl = baseUrl.endsWith("/")
                ? baseUrl.substring(0, baseUrl.length() - 1)
                : baseUrl;
        if (cleanBaseUrl.startsWith("https://")) {
            return "wss://" + cleanBaseUrl.substring("https://".length()) + "/ws/chat";
        }
        if (cleanBaseUrl.startsWith("http://")) {
            return "ws://" + cleanBaseUrl.substring("http://".length()) + "/ws/chat";
        }
        return cleanBaseUrl + "/ws/chat";
    }

    // Helper to check if the current user is a customer
    private boolean isCustomer() {
        return "CUSTOMER".equals(role);
    }

    // if connection drops, try to reconnect after 1 second
    private void scheduleReconnect() {
        if (manualDisconnect) {
            return;
        }
        reconnectHandler.removeCallbacksAndMessages(null);
        reconnectHandler.postDelayed(this::connect, 1000);
    }
}