package io.trygvis.rules.wireguard;

import java.util.ArrayList
import io.trygvis.rules.dns.DnsEntry;
import io.trygvis.rules.machine.Machine;
import io.trygvis.rules.network.Ipv4Address
import io.trygvis.rules.network.Ipv4Cidr
import io.trygvis.rules.network.IpCalc
import java.util.Map

global io.trygvis.rules.engine.TemplateEngine te;

dialect  "mvel"

declare WgNet
    name        : String
    domain      : String
    linkCidr    : String
    networkCidr : String
    networkBits : int
end

declare WgIpPool
    net  : String
    role : String
    cidr : Ipv4Cidr
end

declare WgHost
    name       : String
    net        : String
    publicName : String
    netToNetIp : String
    networkIp  : String
end

declare WgConnection
    host : String
    to   : String
end

declare WgIpAllocation
    host : String
    role : String
    ip   : Ipv4Address
end

declare WgNetworkAllocation
    host : String
    role : String
    cidr : Ipv4Cidr
end

rule "Create IP pools" when
    $net : WgNet()
//    not(Ipv4Cidr(network == Ipv4Cidr.parseCidr($net.linkCidr).network))
then
    System.out.println("Creating main IP pools");
    insert(new WgIpPool($net.name, "link", Ipv4Cidr.parseCidr($net.linkCidr)))
    insert(new WgIpPool($net.name, "networks", Ipv4Cidr.parseCidr($net.networkCidr)))
end

rule "WgHost VPN machines"
when
    $machine : Machine()
    $wgNet : WgNet(name == "vpn0")
    not(WgHost(name == $machine.name))
then
    var wgHost = new WgHost();
    wgHost.name = $machine.name;
    wgHost.net = $wgNet.name;
    wgHost.publicName = $machine.fqdn;
    insert(wgHost)
end

rule "Set public name of WgHost"
when
    $host : WgHost(publicName == null)
    $m    : Machine(name == $host.name, fqdn != null)
then
    modify($host) {
        publicName = $m.fqdn
    }
end

rule "Make DNS entries for all VPN hosts"
when
    $h : WgHost()
    $net : WgNet(name == $h.net)
    not(DnsEntry(fqdn == "%s.%s".formatted($h.name, $net.domain), type == "A"))
then
    var fqdn = "%s.%s".formatted($h.name, $net.domain);
    insert(DnsEntry.a(fqdn))
end

rule "Connect VPN nodes"
    salience -1
when
    $h : WgHost()
    $other : WgHost(publicName != null, name != $h.name)
then
    System.out.printf("VPN connection from %s to %s%n", $h.name, $other.name);
    insert(new WgConnection($h.name, $other.name))
end

rule "Assign link IP"
when
    $net : WgNet()
    $host : WgHost(net == $net.name)
    $pool : WgIpPool(net == $net.name, role == "link")
    not(WgIpAllocation(host == $host.name, role == $pool.role))
    $ip : Ipv4Address() from $pool.cidr.addresses()
    not(WgIpAllocation(ip == $ip))
then
    System.out.printf("IP: net=%s, pool.role=%s, host=%s, ip=%s%n", $net.name, $pool.role, $host.name, $ip);
    insert(new WgIpAllocation($host.name, $pool.role, $ip))
end

rule "Assign network CIDR"
when
    $net : WgNet()
    $host : WgHost(net == $net.name)
    $network : Ipv4Cidr() from Ipv4Cidr.parseCidr($net.networkCidr).partition($net.networkBits)
    not(WgNetworkAllocation(host == $host.name, role == "network"))
    not(WgNetworkAllocation(cidr == $network))
then
    System.out.printf("Network CIDR: net=%s, host=%s, network=%s%n", $net.name, $host.name, $network);
    insert(new WgNetworkAllocation($host.name, "network", $network))
end

rule "Generate per-net files"
    agenda-group "generate"
    salience 10
when
    $net : WgNet()
    $names : ArrayList() from accumulate(WgHost(net == $net.name, $name: name), collectList($name))
    $hosts : ArrayList() from accumulate(Machine($names contains name, $m: this), collectList($m))
then
    te.template("wireguard/ansible", "wireguard-" + $net.name + ".yml", Map.of(
            "net", $net
    ));

    te.template("wireguard/inventory", "inventory.yml", Map.of(
            "hosts", $hosts
    ));
end

rule "Generate per-net, per-host files"
    agenda-group "generate"
    salience 10
when
    $net : WgNet()
    $host : WgHost(net == $net.name)
    $link : WgIpAllocation(host == $host.name, role == "link")
    $network : WgNetworkAllocation(host == $host.name, role == "network")
    $peerMachines : ArrayList() from accumulate(WgConnection(host == $host.name, $to: to), collectList($to))
    $peers : ArrayList() from accumulate(Machine($peerMachines contains name, $fqdn: fqdn), collectList($fqdn))
then
    System.out.printf("Generating per-host files: net=%s, host=%s%n", $net.name, $host.name);

    String output = "host_vars/%s/wireguard.yml".formatted($host.name);

    te.template("wireguard/ansible-host", output, Map.of(
            "net", $net,
            "host", $host,
            "link", $link.ip,
            "network", $network.cidr,
            "peers", $peers
    ));
end