package sync.source.rpc;

import sync.model.RecoveryResult;
import sync.model.SnapshotResult;
import sync.source.AccountUpdateListener;
import sync.source.ConnectionListener;
import sync.source.SolanaDataSource;

import java.io.IOException;
import java.util.List;
import java.util.Set;

public final class RpcSolanaDataSource
        implements SolanaDataSource {

    private final SolanaRpcClient
            rpcClient;

    private final SolanaWebSocketClient
            webSocketClient;

    public RpcSolanaDataSource(
            String rpcUrl,
            String websocketUrl,
            String programId,
            String commitment
    ) {

        this.rpcClient =
                new SolanaRpcClient(
                        rpcUrl,
                        programId,
                        commitment
                );

        this.webSocketClient =
                new SolanaWebSocketClient(
                        websocketUrl,
                        programId,
                        commitment
                );
    }

    @Override
    public SnapshotResult loadFullSnapshot()
            throws Exception {

        return rpcClient
                .loadFullSnapshot();
    }

    @Override
    public void startRealtime(
            AccountUpdateListener accountUpdateListener,
            ConnectionListener connectionListener
    ) {

        webSocketClient.start(
                accountUpdateListener,
                connectionListener
        );
    }

    @Override
    public void stopRealtime() {

        webSocketClient.stop();
    }

    @Override
    public RecoveryResult recoverFrom(
            long lastReconciledSlot
    ) throws Exception {

        /*
         * Fix the upper boundary at the beginning.
         *
         * Realtime is already active and will handle
         * anything newer than targetSlot.
         */
        long targetSlot =
                rpcClient
                        .getCurrentSlot();

        System.out.printf(
                "Incremental recovery requested. fromSlot=%d targetSlot=%d%n",
                lastReconciledSlot,
                targetSlot
        );

        if (targetSlot
                <= lastReconciledSlot) {

            return RecoveryResult.complete(
                    lastReconciledSlot,
                    targetSlot,
                    List.of(),
                    List.of()
            );
        }

        long firstAvailableBlock =
                rpcClient
                        .getFirstAvailableBlock();

        if (firstAvailableBlock < 0) {

            return RecoveryResult.incomplete(
                    lastReconciledSlot,
                    targetSlot,
                    "cannot_resolve_first_available_block"
            );
        }

        /*
         * The RPC node has already pruned the checkpoint
         * range. Incremental recovery cannot be proven.
         */
        if (lastReconciledSlot
                < firstAvailableBlock) {

            return RecoveryResult.incomplete(
                    lastReconciledSlot,
                    targetSlot,
                    "checkpoint_older_than_rpc_history"
            );
        }

        try {

            List<SolanaRpcClient.SignatureInfo>
                    signatures =
                    rpcClient
                            .getSignaturesAfterSlot(
                                    lastReconciledSlot,
                                    targetSlot
                            );

            System.out.printf(
                    "Recovery transactions found: %d%n",
                    signatures.size()
            );

            if (signatures.isEmpty()) {

                return RecoveryResult.complete(
                        lastReconciledSlot,
                        targetSlot,
                        List.of(),
                        List.of()
                );
            }

            Set<String> touchedAddresses =
                    rpcClient
                            .getTouchedAddresses(
                                    signatures
                            );

            System.out.printf(
                    "Recovery unique touched addresses: %d%n",
                    touchedAddresses.size()
            );

            if (touchedAddresses.isEmpty()) {

                return RecoveryResult.complete(
                        lastReconciledSlot,
                        targetSlot,
                        List.of(),
                        List.of()
                );
            }

            SolanaRpcClient.AccountBatchResult
                    accountBatchResult =
                    rpcClient
                            .getCurrentAccounts(
                                    touchedAddresses,
                                    targetSlot
                            );

            System.out.printf(
                    "Recovery state resolved: updates=%d missing=%d%n",
                    accountBatchResult
                            .updates()
                            .size(),
                    accountBatchResult
                            .missingAddresses()
                            .size()
            );

            return RecoveryResult.complete(
                    lastReconciledSlot,
                    targetSlot,
                    accountBatchResult
                            .updates(),
                    accountBatchResult
                            .missingAddresses()
            );

        } catch (IOException exception) {

            /*
             * Missing transaction history, unsupported
             * transaction format or unavailable RPC history
             * means we cannot claim a complete recovery.
             *
             * Coordinator will fall back to a full snapshot.
             */

            System.err.println(
                    "Incremental recovery RPC failure: "
                            + exception.getMessage()
            );

            return RecoveryResult.incomplete(
                    lastReconciledSlot,
                    targetSlot,
                    "rpc_history_recovery_failed: "
                            + exception.getMessage()
            );
        }
    }

    @Override
    public long getCurrentSlot()
            throws Exception {

        return rpcClient
                .getCurrentSlot();
    }

    @Override
    public void close() {

        webSocketClient.close();

        rpcClient.close();
    }
}