package com.dragonssh.xhttpdemo.core.tunnel.vpn;

import android.net.VpnService;

import java.lang.ref.WeakReference;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import java.net.DatagramSocket;
import java.net.Socket;

/**
 * Centralized helper to ensure control sockets (SSH / proxy / TLS) are excluded
 * from the VPN routing.
 *
 * <p>When a VPN is established and a control socket is not protected, Android
 * will route that socket into the VPN itself, creating a loop (VPN -> SOCKS ->
 * VPN). This is especially visible with IPv6, where ::/0 routing is commonly
 * installed.
 */
public final class SocketProtector {

    private static volatile VpnService sVpnService;

    // True only after the active VPN interface has been established with this
    // application UID excluded through Builder.addDisallowedApplication().
    // In that mode every control-plane socket created by this process already
    // uses the underlying network, even on devices where protect(Socket)
    // incorrectly returns false during a reconnect.
    private static volatile boolean sProcessBypassConfigured;

    // Sockets created before the VpnService instance is available.
    // We keep weak refs to avoid leaking in edge cases.
    private static final Object sLock = new Object();
    private static final List<WeakReference<Object>> sPending = new ArrayList<>();

    private SocketProtector() {
    }

    public static void setVpnService(VpnService vpnService) {
        sVpnService = vpnService;
        if (vpnService != null) {
            flushPending(vpnService);
        }
    }

    /**
     * Clears the active VpnService only if it is still the same owner.
     * This prevents a stopped XRAY service from clearing the SSH VPN service
     * after the user switches modes.
     */
    public static void clearVpnService(VpnService vpnService) {
        if (vpnService == null) {
            return;
        }
        synchronized (sLock) {
            if (sVpnService == vpnService) {
                sVpnService = null;
                sProcessBypassConfigured = false;
                sPending.clear();
            }
        }
    }

    public static boolean isVpnServiceReady() {
        return sVpnService != null;
    }

    public static void setProcessBypassConfigured(boolean configured) {
        sProcessBypassConfigured = configured;
    }

    public static boolean isProcessBypassConfigured() {
        return sProcessBypassConfigured;
    }

    public static boolean protect(Socket socket) {
        if (socket == null) return false;
        final VpnService svc = sVpnService;
        if (svc == null) {
            enqueue(socket);
            return false;
        }
        return doProtect(svc, socket);
    }

    public static boolean protect(DatagramSocket socket) {
        if (socket == null) return false;
        final VpnService svc = sVpnService;
        if (svc == null) {
            enqueue(socket);
            return false;
        }
        return doProtect(svc, socket);
    }

    /**
     * Protect a raw native socket file descriptor.
     *
     * Use this from gomobile/JNI callbacks before the native code connects
     * the socket. Raw FDs cannot be safely queued: the native side may connect
     * or close/reuse the descriptor before the VPN service becomes available.
     */
    public static boolean protectFd(int fd) {
        if (fd < 0) return false;
        final VpnService svc = sVpnService;
        if (svc == null) return false;
        return doProtectFd(svc, fd);
    }

    private static void enqueue(Object socket) {
        synchronized (sLock) {
            // Avoid unlimited growth if protect() is called repeatedly.
            if (sPending.size() > 128) {
                // Drop oldest.
                sPending.remove(0);
            }
            sPending.add(new WeakReference<>(socket));
        }
    }

    private static void flushPending(VpnService svc) {
        final List<Object> toProtect = new ArrayList<>();
        synchronized (sLock) {
            for (WeakReference<Object> ref : sPending) {
                Object obj = ref.get();
                if (obj != null) {
                    toProtect.add(obj);
                }
            }
            sPending.clear();
        }

        for (Object obj : toProtect) {
            if (obj instanceof Socket) {
                doProtect(svc, (Socket) obj);
            } else if (obj instanceof DatagramSocket) {
                doProtect(svc, (DatagramSocket) obj);
            }
        }
    }

    private static boolean doProtect(VpnService svc, Socket socket) {
        try {
            return svc.protect(socket);
        } catch (Throwable ignored) {
            return false;
        }
    }

    private static boolean doProtect(VpnService svc, DatagramSocket socket) {
        try {
            return svc.protect(socket);
        } catch (Throwable ignored) {
            return false;
        }
    }

    private static boolean doProtectFd(VpnService svc, int fd) {
        try {
            return svc.protect(fd);
        } catch (Throwable ignored) {
            return false;
        }
    }
}
