/*
 * SPDX-License-Identifier: Apache-2.0
 */

package com.wireguard.android.backend;

import android.net.TrafficStats;
import android.util.Log;

import com.lumaserv.bgp.BGPListener;
import com.lumaserv.bgp.BGPServer;
import com.lumaserv.bgp.BGPSession;
import com.lumaserv.bgp.BGPSessionConfiguration;
import com.lumaserv.bgp.protocol.AFI;
import com.lumaserv.bgp.protocol.BGPPacket;
import com.lumaserv.bgp.protocol.IPPrefix;
import com.lumaserv.bgp.protocol.SAFI;
import com.lumaserv.bgp.protocol.attribute.ASPathAttribute;
import com.lumaserv.bgp.protocol.attribute.MPReachableNLRIAttribute;
import com.lumaserv.bgp.protocol.attribute.NextHopAttribute;
import com.lumaserv.bgp.protocol.attribute.OriginAttribute;
import com.lumaserv.bgp.protocol.attribute.PathAttribute;
import com.lumaserv.bgp.protocol.attribute.TunnelEncapsAttribute;
import com.lumaserv.bgp.protocol.message.BGPUpdate;

import com.wireguard.android.backend.Backend;
import com.wireguard.config.InetEndpoint;
import com.wireguard.config.InetNetwork;
import com.wireguard.crypto.Key;
import com.wireguard.crypto.KeyFormatException;

import io.grpc.ManagedChannel;

import java.io.IOException;
import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketException;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.List;

import javax.net.SocketFactory;

public class Bgp implements BGPListener {
    private static final String TAG = "WireGuard/Bgp";
    private static final String SESSION = "demosession";
    private static final int MY_ASN = (int)4200000201L;
    private static final int REMOTE_ASN = (int)4200000010L;
    private static final String REMOTE_ADDR = "10.49.32.1";
    private static final String REMOTE_ID = "10.49.160.1";
    private static final String LOCAL_ID = "10.49.33.218";
    private static final int PORT = 0;
    private static final int STATS_TAG = 1; // FIXME

    private final Backend backend;
    private final ManagedChannel channel;
    private final Tunnel tunnel;
    private final int tunnelHandle;
    private BGPServer server;

    public Bgp(Backend backend, ManagedChannel channel, Tunnel tunnel, int tunnelHandle) {
	this.backend = backend;
	this.channel = channel;
	this.tunnel = tunnel;
	this.tunnelHandle = tunnelHandle;
    }

    @Override
    public void onOpen(BGPSession session) {
	// DO WHAT YOU WANT
	Log.i(TAG, "onOpen");

	// BGPUpdate update;
	// {
	//     List<IPPrefix> prefixes = new ArrayList<>(1);
	//     prefixes.add(new IPPrefix(new byte[]{10, 49, 124, 105}, 32));
	//     List<PathAttribute> attrs = new ArrayList<>();
	//     attrs.add(new OriginAttribute(OriginAttribute.Origin.IGP));
	//     attrs.add(new NextHopAttribute().setAddress(new byte[]{10, 49, 125, 105}));
	//     ASPathAttribute.Segment seg = new ASPathAttribute.Segment();
	//     seg.setType(ASPathAttribute.Segment.Type.SEQUENCE);
	//     seg.getAsns().add(MY_ASN);
	//     ASPathAttribute asPath = new ASPathAttribute(session);
	//     asPath.getSegments().add(seg);
	//     attrs.add(asPath);
	//     update = new BGPUpdate().setAttributes(attrs).setPrefixes(prefixes);
	// }

	// BGPUpdate update2;
	// try {
	//     List<PathAttribute> attrs = new ArrayList<>();
	//     attrs.add(new OriginAttribute(OriginAttribute.Origin.IGP));
	//     ASPathAttribute.Segment seg = new ASPathAttribute.Segment();
	//     seg.setType(ASPathAttribute.Segment.Type.SEQUENCE);
	//     seg.getAsns().add(MY_ASN);
	//     ASPathAttribute asPath = new ASPathAttribute(session);
	//     List<IPPrefix> prefixes = new ArrayList<>(1);
	//     prefixes.add(new IPPrefix(new byte[]{0x20, 0x1, 0x04, 0x70, (byte)0xdf, (byte)0xae, 0x63, 0, 0, 0, 0, 0, 0, 0, 0x01, 0x05}, 128));
	//     MPReachableNLRIAttribute mpr = new MPReachableNLRIAttribute();
	//     mpr.setAfi(AFI.IPV6).setSafi(SAFI.UNICAST).setNextHop(InetAddress.getByName("2001:470:dfae:6300::1:105")).setNlriPrefixes(prefixes);
	//     attrs.add(mpr);
	//     asPath.getSegments().add(seg);
	//     attrs.add(asPath);
	//     update2 = new BGPUpdate().setAttributes(attrs);
	// } catch (UnknownHostException ex) {
	//     throw new RuntimeException(ex);
	// }

	// try {
	//     session.sendUpdate(update);
	//     session.sendUpdate(update2);
	// } catch (IOException ex) {
	//     throw new RuntimeException(ex);
	// }
    }

    @Override
    public void onUpdate(BGPSession session, BGPUpdate update) {
	// DO WHAT YOU WANT
	Log.i(TAG, "onUpdate: " + update.getPrefixes() + ",-" + update.getWithdrawnPrefixes() + "," + update.getAttributes());
	MPReachableNLRIAttribute mpr = null;
	TunnelEncapsAttribute te = null;

	for (PathAttribute attr: update.getAttributes()) {
	    if (attr instanceof TunnelEncapsAttribute) {
		te = (TunnelEncapsAttribute)attr;
	    } else if (attr instanceof MPReachableNLRIAttribute) {
		mpr = (MPReachableNLRIAttribute)attr;
	    }
	}

	if (te == null) {
	    return;
	}

	TunnelEncapsAttribute.WireGuard wg = null;
	TunnelEncapsAttribute.Color col = null;
	TunnelEncapsAttribute.EgressEndpoint ep = null;
	TunnelEncapsAttribute.UDPDestinationPort port = null;

	for (TunnelEncapsAttribute.Tunnel t: te.getTunnels()) {
	    if (t.getType() != 51820) {
		continue;
	    }

	    for (TunnelEncapsAttribute.SubTlv st: t.getSubTlvs()) {
		if (st instanceof TunnelEncapsAttribute.WireGuard) {
		    wg = (TunnelEncapsAttribute.WireGuard)st;
		} else if (st instanceof TunnelEncapsAttribute.Color) {
		    col = (TunnelEncapsAttribute.Color)st;
		} else if (st instanceof TunnelEncapsAttribute.EgressEndpoint) {
		    ep = (TunnelEncapsAttribute.EgressEndpoint)st;
		} else if (st instanceof TunnelEncapsAttribute.UDPDestinationPort) {
		    port = (TunnelEncapsAttribute.UDPDestinationPort)st;
		}
	    }
	}

	if (wg == null) {
	    return;
	}

	try {
	    Key publicKey = Key.fromBytes(wg.getPublicKey());
	    InetEndpoint endpoint = null;

	    if (ep != null && port != null) {
		endpoint = InetEndpoint.fromAddress(ep.getAddress(), port.getPort());
	    }

	    tunnel.onEndpointChange(publicKey, endpoint);

	    List<InetNetwork> addNetworks = new ArrayList<>();
	    List<InetNetwork> removeNetworks = new ArrayList<>();

	    for (IPPrefix prefix: update.getPrefixes()) {
		try {
		    addNetworks.add(new InetNetwork(InetAddress.getByAddress(prefix.getAddress()), prefix.getLength()));
		} catch (UnknownHostException ignore) {
		}
	    }

	    for (IPPrefix prefix: update.getWithdrawnPrefixes()) {
		try {
		    removeNetworks.add(new InetNetwork(InetAddress.getByAddress(prefix.getAddress()), prefix.getLength()));
		} catch (UnknownHostException ignore) {
		}
	    }

	    if (mpr != null && (mpr.getAfi() == AFI.IPV6 || mpr.getAfi() == AFI.IPV4) && mpr.getSafi() == SAFI.UNICAST) {
		for (IPPrefix prefix: mpr.getNlriPrefixes()) {
		    try {
			addNetworks.add(new InetNetwork(InetAddress.getByAddress(prefix.getAddress()), prefix.getLength()));
		    } catch (UnknownHostException ignore) {
		    }
		}
	    }

	    tunnel.onAllowedIpsChange(publicKey, addNetworks, removeNetworks);
	    // backend.addAllowedIps(tunnel, publicKey, addNetworks);
	    // backend.removeAllowedIps(tunnel, publicKey, addNetworks); // TODO
	} catch (KeyFormatException ex) {
	    Log.w(TAG, "Key.fromBytes " + ex);
	}
    }

    @Override
    public void onClose(BGPSession session) {
	// NOT YET IMPLEMENTED
	Log.i(TAG, "onClose");
    }

    public boolean startServer() {
	stopServer();
	try {
	    SocketFactory factory = new SocketFactory() {
		    private Socket taggedSocket(Socket sock) throws SocketException {
			TrafficStats.tagSocket(sock);
			return sock;
		    }

		    @Override
		    public Socket createSocket(String host, int port) throws IOException {
			TrafficStats.setThreadStatsTag(STATS_TAG);
			return taggedSocket(new Socket(host, port));
		    }

		    @Override
		    public Socket createSocket(InetAddress host, int port) throws IOException {
			TrafficStats.setThreadStatsTag(STATS_TAG);
			return taggedSocket(new Socket(host, port));
		    }

		    @Override
		    public Socket createSocket(InetAddress address, int port, InetAddress localAddress, int localPort) throws IOException {
			TrafficStats.setThreadStatsTag(STATS_TAG);
			return taggedSocket(new Socket(address, port, localAddress, localPort));
		    }

		    @Override
		    public Socket createSocket(String host, int port, InetAddress localHost, int localPort) throws IOException {
			TrafficStats.setThreadStatsTag(STATS_TAG);
			return taggedSocket(new Socket(host, port, localHost, localPort));
		    }
		};

	    BGPSessionConfiguration config =
		new BGPSessionConfiguration(SESSION,
					    MY_ASN,
					    ip(LOCAL_ID),
					    REMOTE_ASN,
					    ip(REMOTE_ID),
					    null, // Remote address
					    factory,
					    this);
	    TrafficStats.setThreadStatsTag(STATS_TAG);
	    ServerSocket socket = new ServerSocket(PORT);
	    //TrafficStats.tagSocket(socket);
	    // Set
	    server = new BGPServer(socket);
	    server.getSessionConfigurations().add(config);
	    server.connect(config, REMOTE_ADDR);
	    return true;
	} catch (IOException ex) {
	    return false;
	}
    }

    public void stopServer() {
	if (server != null) {
	    server.shutdown();
	    server = null;
	}
    }

    private static byte[] ip(String s) throws UnknownHostException {
	InetAddress addr = InetAddress.getByName(s);
	byte[] data = addr.getAddress();
	if (data.length != 4)
	    throw new UnknownHostException(s + ": Not an IPv4 address");
	return data;
    }
}