/*
 * Copyright © 2023 WireGuard LLC. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

package com.wireguard.util;

import java.io.IOException;
import java.net.DatagramSocket;
import java.net.Inet4Address;
import java.net.Inet6Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketException;
import java.net.UnknownHostException;

import android.net.IpPrefix;
import android.net.LinkProperties;
import android.net.Network;
import android.net.TrafficStats;
import android.util.Log;

import androidx.annotation.Nullable;

@NonNullForAll
public class Resolver {
    private static final String TAG = "WireGuard/Resolver";
    private static final int STATS_TAG = 3; // FIXME
    @Nullable private final Network network;
    @Nullable private final LinkProperties linkProps;
    @Nullable private IpPrefix nat64Prefix;

    public Resolver(Network network, LinkProperties linkProps) {
        this.network = network;
        this.linkProps = linkProps;
        if (linkProps != null) {
            this.nat64Prefix = linkProps.getNat64Prefix();
        }
    }

    public static boolean isULA(Inet6Address addr) {
        byte[] raw = addr.getAddress();
        return ((raw[0] & 0xfe) == 0xfc);
    }

    boolean isWithinNAT64Prefix(Inet6Address address) {
        if (nat64Prefix == null)
            return false;

        int prefixLength = nat64Prefix.getPrefixLength();
        byte[] rawAddr = address.getAddress();
        byte[] rawPrefix = nat64Prefix.getRawAddress();

        for (int i=0; i < prefixLength/8; i++) {
            if (rawAddr[i] != rawPrefix[i])
                return false;
        }

        return true;
    }

    boolean isPreferredIPv6(Inet6Address local, Inet6Address remote) {
        if (linkProps == null) {
            // Prefer IPv4 if there are not link properties that can
            // be tested.
            return false;
        }

        // * Prefer IPv4 if local or remote address is ULA
        // * Prefer IPv4 if remote IPv6 is within NAT64 prefix.
        // * Otherwise prefer IPv6
        boolean isLocalULA = isULA(local);
        boolean isRemoteULA = isULA(remote);

        if (isLocalULA || isRemoteULA) {
            return false;
        }

        if (isWithinNAT64Prefix(remote)) {
            return false;
        }

        return true;
    }

    public InetAddress resolve(String host) throws UnknownHostException {
        TrafficStats.setThreadStatsTag(STATS_TAG);
        final InetAddress[] candidates = network != null ? network.getAllByName(host) : InetAddress.getAllByName(host);
        InetAddress address = candidates[0];
        for (final InetAddress candidate : candidates) {
            DatagramSocket sock;

            try {
                sock = new DatagramSocket();
                TrafficStats.tagDatagramSocket(sock);
                if (network != null) {
                    network.bindSocket(sock);
                }
            } catch (SocketException e) {
                // Return first candidate as fallback
                Log.w(TAG, "DatagramSocket failed, fallback to: \"" + address);
                return address;
            } catch (IOException e) {
                // Return first candidate as fallback
                Log.w(TAG, "BindSocket failed, fallback to: \"" + address);
                return address;
            }

            sock.connect(candidate, 51820);

            if (sock.getLocalAddress().isAnyLocalAddress()) {
                // Connect didn't find a local address.
                Log.w(TAG, "No local address");
                continue;
            }

            Log.w(TAG, "Local address: " + sock.getLocalAddress());

            if (candidate instanceof Inet4Address) {
                // Accept IPv4 as preferred address.
                address = candidate;
                break;
            }

            Inet6Address local = (Inet6Address)sock.getLocalAddress();
            InetSocketAddress remoteSockAddr = (InetSocketAddress)sock.getRemoteSocketAddress();
            Inet6Address remote = (Inet6Address)remoteSockAddr.getAddress();
            sock.close();

            if (isPreferredIPv6(local, remote)) {
                address = candidate;
                break;
            }
        }
        Log.w(TAG, "Resolved \"" + host + "\" to: " + address);
        return address;
    }
}