#!/usr/bin/perl

use v5.12;

use strict;
use warnings;

use IO::File;
use Math::Trig qw(deg2rad pip2 great_circle_distance);
use JSON::XS;
use LWP::UserAgent;
use YAML;

my $originfile = shift @ARGV;
my $clusterfile = shift @ARGV;
my $zone = shift @ARGV;
my $jszone = shift @ARGV;
my $targetoriginfile = shift @ARGV;
my $origins = YAML::LoadFile($originfile);
my $clusters = YAML::LoadFile($clusterfile);
my $gdnsname = shift @ARGV;
my @servers;

# Initialise cluster details
while (my($name,$cluster) = each %$clusters)
{
    if ($cluster->{servers})
    {
        $cluster->{bandwidth} = 0;

        foreach my $server (@{$cluster->{servers}})
        {
            $server->{cluster} = $cluster;
            $cluster->{bandwidth} = $cluster->{bandwidth} + $server->{bandwidth};

            push @servers, $server;
        }
    }
    else
    {
        my $server = {
            cluster => $cluster,
            statuscake => $cluster->{statuscake},
            bandwidth => $cluster->{bandwidth},
            ipv4 => $cluster->{ipv4},
            ipv6 => $cluster->{ipv6}
        };

        $cluster->{servers} = [ $server ];

        push @servers, $server;
    }

    $cluster->{name} = $name;
    $cluster->{status} = "down";
}

# Initialise server details
foreach my $server (@servers)
{
    $server->{status} = "up";
}

# If statuscake support is enabled then check which servers are up
if ($ENV{STATUSCAKE_USERNAME} && $ENV{STATUSCAKE_APIKEY})
{
    my $ua = LWP::UserAgent->new;
    my $cache;

    $ua->agent("mkgeo/1.0");
    $ua->default_header("Username", $ENV{STATUSCAKE_USERNAME});
    $ua->default_header("API", $ENV{STATUSCAKE_APIKEY});

    if (-f "statuscake.yml")
    {
        $cache = YAML::LoadFile("statuscake.yml");
    }
    else
    {
        $cache = {};
    }

    my $response = $ua->get("https://app.statuscake.com/API/Tests/");

    if ($response->is_success)
    {
        my $tests = decode_json($response->content);

        foreach my $test (@$tests)
        {
            my $testid = $test->{TestID};

            if ($test->{Status} eq "Up" && !$test->{Paused})
            {
                $cache->{$testid} = "up";
            }
            else
            {
                $cache->{$testid} = "down";
            }
        }
    }

    foreach my $server (@servers)
    {
        if (my $testids = $server->{statuscake})
        {
            $server->{status} = "up";

            for my $testid (@$testids)
            {
                my $testresult = $cache->{$testid} || "down";

                $server->{status} = "down" if $testresult eq "down";
            }
        }
        else
        {
            $server->{status} = "down";
        }
    }

    YAML::DumpFile("statuscake.yml", $cache);
}

# Mark a cluster as up if any servers are up
foreach my $server (@servers)
{
    if ($server->{status} eq "up")
    {
        $server->{cluster}->{status} = "up";
    }
    else
    {
        $server->{cluster}->{bandwidth} = $server->{cluster}->{bandwidth} - $server->{bandwidth};
    }
}

# Create target origins object
my $targetorigins = {};

# Initialise cluster details
while (my($name,$cluster) = each %$clusters)
{
    $cluster->{bandwidth_limit} = $cluster->{bandwidth} * 1024 * 1024;
    $cluster->{bandwidth_used} = 0;

    $targetorigins->{$cluster->{name}} = {
        code => $cluster->{name},
        name => $cluster->{name},
        lat => $cluster->{lat},
        lon => $cluster->{lon},
        bandwidth => 0
    };
}

my @mappings = ();

# Scan origins and work out which clusters each can use
foreach my $origin (values %$origins)
{
    foreach my $cluster (values %$clusters)
    {
        my $match = match_origin($cluster, $origin);

        if ($cluster->{status} eq "up" && $match ne "denied")
        {
            my $priority = $match eq "preferred" ? 20 : 10;
            my $distance = distance($origin->{lat}, $origin->{lon}, $cluster->{lat}, $cluster->{lon});

            push @mappings, {
                origin => $origin, cluster => $cluster,
                priority => $priority, distance => $distance
            };
        }
    }
}

# Allocate each country to a cluster
allocate_clusters(@mappings);

# If we failed to allocate every origin then loop, increasing
# the bandwidth for each cluster by a little and retrying until
# we manage to allocate everything
while (grep { !exists($_->{cluster}) } values %$origins)
{
    # Clear any existing mappings of countries to clusters
    foreach my $origin (values %$origins)
    {
        delete $origin->{cluster};
    }

    # Reset bandwidth usage for clusters and increase limits by 10%
    foreach my $cluster (values %$clusters)
    {
        $cluster->{bandwidth_used} = 0;
        $cluster->{bandwidth_limit} = $cluster->{bandwidth_limit} * 1.1;
    }

    # Try the allocate again
    allocate_clusters(@mappings);
}

# Create JSON collection object
my @json;

# Open output files
my $zonefile = IO::File->new("> data/${zone}") || die "$!";
my $jszonefile = IO::File->new("> include/${jszone}.js") || die "$!";
my $jsonfile = IO::File->new("> json/${zone}.json") || die "$!";

# Output headers
$jszonefile->print("var \U${jszone}\E_RECORDS = [\n");

# Skip outputting server map if gdnsd is doing that
if (!defined($gdnsname))
{
    # Output details for each country
    foreach my $origin (sort { $a->{name} cmp $b->{name} } values %$origins)
    {
        my $cluster = $origin->{cluster};
        my $clon = $origin->{lon};
        my $clat = $origin->{lat};
        my $slon = $cluster->{lon};
        my $slat = $cluster->{lat};

        if ($clon > 0 && $slon < 0 && 360 + $slon - $clon < $clon - $slon)
        {
            $slon = $slon + 360;
        }
        elsif ($slon > 0 && $clon < 0 && 360 + $clon - $slon < $slon - $clon)
        {
            $clon = $clon + 360;
        }

        $zonefile->print("# $origin->{name}\n");
        $zonefile->print("C\L$origin->{code}\E.${zone}:$cluster->{name}.${zone}:600\n");

        $jszonefile->print("  CNAME(\"\L$origin->{code}\E.${jszone}\", \"$cluster->{name}.${zone}.\", TTL(\"10m\")),\n");

        push @json, {
            type => "Feature",
            geometry => {
                type => "LineString",
                coordinates => [ [ $clon, $clat ], [ $slon, $slat ] ]
            },
            properties => {
                origin => $origin->{name},
                server => $cluster->{name},
                colour => $cluster->{colour}
            }
        };

        $targetorigins->{$cluster->{name}}->{bandwidth} += $origin->{bandwidth};
    }

    # Header for default records
    $zonefile->print("# Unknown origins\n");

    # Output default records for IPs that can't be mapped to a country
    foreach my $cluster (sort { $a->{name} cmp $b->{name} } values %$clusters)
    {
        my $name = $cluster->{name};

        if (my $default = $cluster->{default})
        {
            output_server($zonefile, $jszonefile, "${default}.${zone}", "${default}.${jszone}", $cluster);
        }
        elsif (exists($cluster->{default}))
        {
            output_server($zonefile, $jszonefile, "${zone}", "${jszone}", $cluster);
        }
    }
}

# Header for underlying servers
$zonefile->print("# Servers\n");

# Output A records for each cluster
foreach my $cluster (sort { $a->{name} cmp $b->{name} } values %$clusters)
{
    my $name = $cluster->{name};

    output_server($zonefile, $jszonefile, "${name}.${zone}", "${name}.${jszone}", $cluster);

    if (@{$cluster->{servers}} > 1)
    {
        output_server($zonefile, $jszonefile, "${name}-%02d.${zone}", "${name}-%02d.${jszone}", $cluster);
    }
}

# Output the GeoJSON text
$jsonfile->print(encode_json(\@json));

# Output footers
$jszonefile->print("];\n");

# Close the output files
$jsonfile->close();
$jszonefile->close();
$zonefile->close();

# Output gdnsd configuration
if (defined($gdnsname))
{
    my $gdnsmapfile = IO::File->new("> gdns/${gdnsname}.map") || die "$!";
    my $gdnsresourcefile = IO::File->new("> gdns/${gdnsname}.resource") || die "$!";
    my $gdnsweightedfile = IO::File->new("> gdns/${gdnsname}.weighted") || die "$!";
    my $continent = "";

    $gdnsmapfile->print("${gdnsname} => {\n");
    $gdnsmapfile->print("  geoip2_db => /usr/share/GeoIP/GeoLite2-Country.mmdb\n");
    $gdnsmapfile->print("  datacenters => [" . join(",", sort(keys(%$clusters))) . "]\n");
    $gdnsmapfile->print("  map => {\n");
    $gdnsmapfile->print("    default => [" . join(",", sort(map { $_->{name} } grep { $_->{default} } values(%$clusters))) . "]\n");

    foreach my $origin (sort { $a->{continent} cmp $b->{continent} || $a->{code} cmp $b->{code} } values %$origins)
    {
        my $code = $origin->{code};
        my $cluster = $origin->{cluster}->{name};

        next if $code eq "XK";

        if ($continent ne $origin->{continent})
        {
            $gdnsmapfile->print("    }\n") if $continent;

            $continent = $origin->{continent};

            $gdnsmapfile->print("    ${continent} => {\n");
        }

        $gdnsmapfile->print("      ${code} => [${cluster}]\n");
    }

    $gdnsmapfile->print("    }\n") if $continent;

    $gdnsmapfile->print("  }\n");
    $gdnsmapfile->print("}\n");

    $gdnsresourcefile->print("${gdnsname} => {\n");
    $gdnsresourcefile->print("  map => ${gdnsname}\n");
    $gdnsresourcefile->print("  dcmap => {\n");

    foreach my $cluster (sort { $a->{name} cmp $b->{name} } values %$clusters)
    {
        my $name = $cluster->{name};

        if (@{$cluster->{servers}} > 1)
        {
            $gdnsweightedfile->print("${name} => {\n");

            while (my($index,$server) = each @{$cluster->{servers}})
            {
                if ($server->{status} eq "up")
                {
                    my $number = sprintf("%02d", $index + 1);
                    my $bandwidth = $server->{bandwidth};

                    $gdnsweightedfile->print("  ${name}-${number} = [ ${name}-${number}.${zone}., ${bandwidth} ]\n");
                }
            }

            $gdnsweightedfile->print("}\n");

            $gdnsresourcefile->print("    ${name} => %weighted!${name}\n");
        }
        else
        {
            $gdnsresourcefile->print("    ${name} => ${name}.${zone}.\n");
        }
    }

    $gdnsresourcefile->print("  }\n");
    $gdnsresourcefile->print("}\n");

    $gdnsweightedfile->close();
    $gdnsresourcefile->close();
    $gdnsmapfile->close();
}

# Output the target details in origin format if required
YAML::DumpFile($targetoriginfile, $targetorigins) if $targetoriginfile;

exit 0;

#
# Match an origin against a cluster
#
sub match_origin
{
    my $cluster = shift;
    my $origin = shift;
    my $match;

    if ($cluster->{preferred} &&
        $cluster->{preferred}->{origins} &&
        grep { $_ eq $origin->{name} } @{$cluster->{preferred}->{origins}})
    {
        $match = "preferred";
    }
    elsif ($cluster->{allowed} &&
           $cluster->{allowed}->{origins} &&
           grep { $_ eq $origin->{name} } @{$cluster->{allowed}->{origins}})
    {
        $match = "allowed";
    }
    elsif ($cluster->{preferred} &&
           $cluster->{preferred}->{countries} &&
           grep { $_ eq $origin->{country} } @{$cluster->{preferred}->{countries}})
    {
        $match = "preferred";
    }
    elsif ($cluster->{allowed} &&
           $cluster->{allowed}->{countries} &&
           grep { $_ eq $origin->{country} } @{$cluster->{allowed}->{countries}})
    {
        $match = "allowed";
    }
    elsif ($cluster->{denied} &&
           $cluster->{denied}->{countries} &&
           grep { $_ eq $origin->{country} } @{$cluster->{denied}->{countries}})
    {
        $match = "denied";
    }
    elsif ($cluster->{preferred} &&
           $cluster->{preferred}->{continents} &&
           grep { $_ eq $origin->{continent} } @{$cluster->{preferred}->{continents}})
    {
        $match = "preferred";
    }
    elsif ($cluster->{allowed} &&
           $cluster->{allowed}->{continents} &&
           grep { $_ eq $origin->{continent} } @{$cluster->{allowed}->{continents}})
    {
        $match = "allowed";
    }
    elsif ($cluster->{denied} &&
           $cluster->{denied}->{continents} &&
           grep { $_ eq $origin->{continent} } @{$cluster->{denied}->{continents}})
    {
        $match = "denied";
    }
    elsif ($cluster->{allowed})
    {
        $match = "denied";
    }
    else
    {
        $match = "allowed";
    }

    return $match;
}

#
# Compute the great circle distance between two points
#
sub distance
{
    my $lat1 = deg2rad(shift);
    my $lon1 = deg2rad(shift);
    my $lat2 = deg2rad(shift);
    my $lon2 = deg2rad(shift);

    return great_circle_distance($lon1, pip2 - $lat1, $lon2, pip2 - $lat2);
}

#
# Allocate each origin to a cluster
#
sub allocate_clusters
{
    my @mappings = sort { compare_mappings($a, $b) } @_;

    # Loop over the mappings, trying to assign each origin to the
    # nearest cluster, but subject to the bandwidth limits
    while (my $mapping = shift @mappings)
    {
        my @group;

        push @group, $mapping;

        while (@mappings && compare_mappings($mapping, $mappings[0]) == 0)
        {
            push @group, shift @mappings;
        }

        for my $mapping (sort compare_bandwidth @group)
        {
            my $origin = $mapping->{origin};
            my $cluster = $mapping->{cluster};

            if (!exists($origin->{cluster}) &&
                $cluster->{bandwidth_used} + $origin->{bandwidth} <= $cluster->{bandwidth_limit})
            {
                $origin->{cluster} = $cluster;
                $cluster->{bandwidth_used} = $cluster->{bandwidth_used} + $origin->{bandwidth};
            }
        }
    }

    return;
}

#
# Compare two mappings to decide which to use
#
sub compare_mappings
{
    my $a = shift;
    my $b = shift;

    return $b->{priority} <=> $a->{priority} ||
           $a->{distance} <=> $b->{distance};
}

#
# Compare two mappings to decide which to try first
#
sub compare_bandwidth
{
    my $a_used = ( $a->{cluster}->{bandwidth_used} * 100.0 ) / ( $a->{cluster}->{bandwidth_limit} * 1.0 );
    my $b_used = ( $b->{cluster}->{bandwidth_used} * 100.0 ) / ( $b->{cluster}->{bandwidth_limit} * 1.0 );

    return $a_used <=> $b_used;
}

#
# Output DNS records for a server
#
sub output_server
{
    my $zonefile = shift;
    my $jszonefile = shift;
    my $name = shift;
    my $jsname = shift;
    my $cluster = shift;

    while (my($index,$server) = each @{$cluster->{servers}})
    {
        if ($server->{status} eq "up")
        {
            $zonefile->printf("+${name}:$server->{ipv4}:600\n", $index + 1);
            $jszonefile->printf("  A(\"${jsname}\", \"$server->{ipv4}\", TTL(\"10m\")),\n", $index + 1);

            if ($server->{ipv6})
            {
                my $ipv6 = $server->{ipv6};

                $ipv6 =~ s/([0-9a-f]{4})(?=.)/$1:/ig;

                $zonefile->printf("3${name}:$server->{ipv6}:600\n", $index + 1);
                $jszonefile->printf("  AAAA(\"${jsname}\", \"${ipv6}\", TTL(\"10m\")),\n", $index + 1);
            }
        }
    }

    return;
}