#!/usr/bin/perl

use strict;
use warnings;

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

my $source = shift @ARGV;
my $zone = shift @ARGV;
my $clusters = YAML::LoadFile("src/${source}");
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,
            pingdom => $cluster->{pingdom},
            bandwidth => $cluster->{bandwidth},
            ipv4 => $cluster->{ipv4},
            ipv6 => $cluster->{ipv6}
        };

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

        push @servers, $server;
    }

    $cluster->{name} = $name;

    if ($ENV{PINGDOM_USERNAME} && $ENV{PINGDOM_PASSWORD})
    {
        $cluster->{status} = "down";
    }
    else
    {
        $cluster->{status} = "up";
    }
}

# Initialise server details
foreach my $server (@servers)
{
    if ($ENV{PINGDOM_USERNAME} && $ENV{PINGDOM_PASSWORD})
    {
        $server->{status} = "down";
    }
    else
    {
        $server->{status} = "up";
    }
}

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

    $ua->default_header("App-Key", "2cohi62u5haxvqmypk3ljqqrze1jufrh");
    $ua->credentials("api.pingdom.com:443", "Pingdom API", $ENV{PINGDOM_USERNAME}, $ENV{PINGDOM_PASSWORD});

    foreach my $server (@servers)
    {
        if (my $checkid = $server->{pingdom})
        {
            my $response = $ua->get("https://api.pingdom.com/api/2.0/checks/${checkid}");

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

                $server->{status} = $check->{check}->{status};

                if ($server->{status} eq "up")
                {
                    $server->{cluster}->{status} = "up";
                }
                else
                {
                    $server->{cluster}->{bandwidth} = $server->{cluster}->{bandwidth} - $server->{bandwidth};
                }
            }
        }
    }
}

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

my %countries = ();
my @mappings = ();

# Create a parser for the country database
my $countries = XML::TreeBuilder->new;

# Parse the country database
$countries->parsefile("lib/countries.xml");

# Load the per-country bandwidth details
my $bandwidth = YAML::LoadFile("bandwidth/${source}.yml");

# Fill in country table and work out which clusters each can use
foreach my $country ($countries->look_down("_tag" => "country"))
{
    my $code = $country->look_down("_tag" => "countryCode")->as_text;
    my $name = $country->look_down("_tag" => "countryName")->as_text;
    my $population = $country->look_down("_tag" => "population")->as_text;
    my $bandwidth = $bandwidth->{$code} || 0;
    my $continent = $country->look_down("_tag" => "continent")->as_text;
    my $west = $country->look_down("_tag" => "west")->as_text;
    my $north = $country->look_down("_tag" => "north")->as_text;
    my $east = $country->look_down("_tag" => "east")->as_text;
    my $south = $country->look_down("_tag" => "south")->as_text;
    my $lat = centre_lat( $south, $north );
    my $lon = centre_lon( $west, $east );

    $countries{$code} = {
        code => $code, name => $name, continent => $continent,
        bandwidth => $bandwidth, lat => $lat, lon => $lon
    };

    foreach my $cluster (values %$clusters)
    {
        my $match = match_country($cluster, $code, $continent);

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

            push @mappings, {
                country => $countries{$code}, cluster => $cluster,
                priority => $priority, distance => $distance
            };
        }
    }
}

# Discard the parsed country database
$countries->delete;

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

# If we failed to allocate every country then loop, increasing
# the bandwidth for each cluster by a little and retrying until
# we manage to allocate everything
while (grep { !exists($_->{cluster}) } values %countries)
{
    # Clear any existing mappings of countries to clusters
    foreach my $country (values %countries)
    {
        delete $country->{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 $jsonfile = IO::File->new("> json/${zone}.json") || die "$!";

# Output details for each country
foreach my $country (values %countries)
{
    my $cluster = $country->{cluster};
    my $clon = $country->{lon};
    my $clat = $country->{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("# $country->{name}\n");
    $zonefile->print("C\L$country->{code}\E.${zone}:$cluster->{name}.${zone}:600\n");

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

# Output default records for IPs that can't be mapped to a country
foreach my $cluster (grep { $clusters->{$_}->{default} } keys %$clusters)
{
    $zonefile->print("# Unknown countries\n");
    $zonefile->print("Cxx.${zone}:${cluster}.${zone}:600\n");
}

$zonefile->print("# Servers\n");

# Output A records for each cluster
while (my($name,$cluster) = each %$clusters)
{
    foreach my $server (@{$cluster->{servers}})
    {
        if ($server->{status} eq "up")
        {
            $zonefile->print("+${name}.${zone}:$server->{ipv4}:600\n");

            if ($server->{ipv6})
            {
#                $zonefile->print("3${name}.${zone}:$server->{ipv6}:600\n");
            }
        }
    }
}

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

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

exit 0;

#
# Find the centre value between two latitudes
#
sub centre_lat
{
    my $south = shift;
    my $north = shift;

    return ( $south + $north ) / 2;
}

#
# Find the centre value between two longitudes
#
sub centre_lon
{
    my $west = shift;
    my $east = shift;
    my $lon;

    if ($west < $east)
    {
        $lon = ( $west + $east ) / 2;
    }
    else
    {
        $lon = ( $west + $east + 360 ) / 2;
    }

    $lon = $lon - 360 if $lon > 180;

    return $lon
}

#
# Match a country against a cluster
#
sub match_country
{
    my $cluster = shift;
    my $country = shift;
    my $continent = shift;
    my $match;

    if ($cluster->{preferred} &&
        $cluster->{preferred}->{countries} &&
        grep { $_ eq $country } @{$cluster->{preferred}->{countries}})
    {
        $match = "preferred";
    }
    elsif ($cluster->{preferred} &&
           $cluster->{preferred}->{continents} &&
           grep { $_ eq $continent } @{$cluster->{preferred}->{continents}})
    {
        $match = "preferred";
    }
    elsif ($cluster->{allowed} &&
           $cluster->{allowed}->{countries} &&
           grep { $_ eq $country } @{$cluster->{allowed}->{countries}})
    {
        $match = "allowed";
    }
    elsif ($cluster->{allowed} &&
           $cluster->{allowed}->{continents} &&
           grep { $_ eq $continent } @{$cluster->{allowed}->{continents}})
    {
        $match = "allowed";
    }
    elsif ($cluster->{denied} &&
           $cluster->{denied}->{countries} &&
           grep { $_ eq $country } @{$cluster->{preferred}->{countries}})
    {
        $match = "denied";
    }
    elsif ($cluster->{denied} &&
           $cluster->{denied}->{continents} &&
           grep { $_ eq $continent } @{$cluster->{preferred}->{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 country to a cluster
#
sub allocate_clusters
{
    my $mappings = shift;

    # Loop over the mappings, trying to assign each country to the
    # nearest cluster, but subject to the bandwidth limits
    foreach my $mapping (sort {  $b->{priority} <=> $a->{priority} || $a->{distance} <=> $b->{distance} } @$mappings)
    {
        my $country = $mapping->{country};
        my $cluster = $mapping->{cluster};

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

    return;
}