#!/usr/bin/perl

use strict;
use warnings;
use Gtk2 -init;
use Cairo;
use Glib;
use POSIX qw(tan);
use constant PI => 3.14159265;

# --- Configuration ---
# This uses the normal connmap configuration dir and resources. But set it to whatever you want.
# You can get the background image needed by $ wget http://superkuh.com/w1000-old.png
# You can get the ipv4 database from https://github.com/h2337/connmap/raw/refs/heads/master/connmap/resources/ipv4.csv.zip

my $home_dir = $ENV{'HOME'} or die "HOME environment variable not set";
my $resource_dir = "$home_dir/.config/connmap/resources";
my $image_path = "$resource_dir/w1000-old.png";
my $ipv4_db_path = "$resource_dir/ipv4.csv";

# --- Global State ---
my $state = {
    current_points => [],
    map_width      => 1000,
    map_height     => 500,
    dot_diameter   => 8,
    dot_color      => [1, 0.2, 0.2, 0.8], # Semi-transparent red
};

#======================================================================
# IP Geolocation and Coordinate Functions (RAM-Efficient)
#======================================================================

sub ip_to_decimal {
    my ($ip_str) = @_;
    my @parts = split /\./, $ip_str;
    return 0 unless @parts == 4;
    return ($parts[0] << 24) + ($parts[1] << 16) + ($parts[2] << 8) + $parts[3];
}

# RAM-EFFICIENT LOOKUP
# Performs a binary search directly on the CSV file on disk.
sub find_location_on_disk {
    my ($fh, $ip_decimal, $file_size) = @_;

    my $low_pos = 0;
    my $high_pos = $file_size;

    # Limit search iterations to prevent infinite loops on malformed files
    for (my $i = 0; $i < 30; $i++) {
        # Stop if the search window is too small
        last if ($high_pos - $low_pos) < 2;

        my $mid_pos = int($low_pos + ($high_pos - $low_pos) / 2);

        # Seek to the middle of the current byte range
        seek($fh, $mid_pos, 0);

        # The seek likely landed us in the middle of a line.
        # Read and discard the partial line to get to a clean start.
        <$fh>;

        # Now, read the next complete line. This is our sample.
        my $line = <$fh>;
        return undef unless defined $line; # End of file or read error
        chomp $line;

        my ($start, $end, $lat, $lon) = split ',', $line;
        # Skip if the line isn't a valid entry
        next unless defined $start and defined $end;

        # Check if we found the correct range
        if ($ip_decimal >= $start && $ip_decimal <= $end) {
            return { lat => $lat, lon => $lon };
        }

        # Adjust the search window (in bytes)
        if ($ip_decimal < $start) {
            $high_pos = $mid_pos;
        } else {
            $low_pos = $mid_pos;
        }
    }

    return undef; # Not found after all iterations
}


sub get_peer_ips {
    my $ss_cmd = q(ss -t -n -4 state established | awk 'NR>1 {print $4}');
    my $output = `$ss_cmd`;
    
    my %unique_ips;
    for my $addr_port (split /\s+/, $output) {
        my ($ip) = ($addr_port =~ /^(.+):\d+$/);
        if ($ip) {
            $unique_ips{$ip} = 1;
        }
    }
    return keys %unique_ips;
}

sub latlon_to_xy {
    my ($latitude, $longitude) = @_;
    my $width = $state->{map_width};
    my $height = $state->{map_height};
    my $xOffset = -29.0;
    my $yOffset = 76.0;
    my $x = ($longitude + 180) * ($width / 360) + $xOffset;
    my $latRad = $latitude * PI / 180;
    my $mercN = log(tan((PI / 4) + ($latRad / 2)));
    my $y = ($height / 2) - ($width * $mercN / (2 * PI)) + $yOffset;
    return ($x, $y);
}

#======================================================================
# GTK, Drawing, and Timer Functions
#======================================================================

sub update_connections_and_redraw {
    my ($drawing_area) = @_;

    # Open the database file for this update cycle.
    open my $db_fh, '<', $ipv4_db_path or do {
        warn "Could not open IP database '$ipv4_db_path': $!";
        return 1; # Keep timer running
    };
    my $db_size = -s $db_fh;

    my @peer_ips = get_peer_ips();
    my @new_points;

    for my $ip (@peer_ips) {
        my $ip_decimal = ip_to_decimal($ip);
        next unless $ip_decimal;
        # Perform the on-disk search for each IP
        my $location = find_location_on_disk($db_fh, $ip_decimal, $db_size);
        if ($location) {
            push @new_points, $location;
        }
    }
    
    close $db_fh; # Close the file handle for this cycle

    $state->{current_points} = \@new_points;
    $drawing_area->queue_draw();
    
    return 1;
}

sub on_draw {
    my ($widget, $event) = @_;
    my $pixbuf = $widget->{_pixbuf_};
    return 1 unless $pixbuf;

    my $cr = Gtk2::Gdk::Cairo::Context->create($widget->window);
    $cr->set_source_pixbuf($pixbuf, 0, 0);
    $cr->paint;
    $cr->set_source_rgba(@{$state->{dot_color}});
    my $radius = $state->{dot_diameter} / 2;

    for my $point (@{$state->{current_points}}) {
        my ($x, $y) = latlon_to_xy($point->{lat}, $point->{lon});
        $cr->arc($x, $y, $radius, 0, 2 * PI);
        $cr->fill;
    }
    
    return 1;
}

#======================================================================
# Main Execution Block
#======================================================================


my $pixbuf;
eval {
    $pixbuf = Gtk2::Gdk::Pixbuf->new_from_file_at_size(
        $image_path, $state->{map_width}, $state->{map_height}
    );
};
die "Error loading image '$image_path': $@" if $@;

my $window = Gtk2::Window->new('toplevel');
$window->set_title("Connection Map (RAM Efficient)");
$window->set_default_size($state->{map_width}, $state->{map_height});
$window->set_position('center');
$window->set_decorated(0);
$window->set_app_paintable(1);

my $screen = $window->get_screen;
my $rgba_colormap = $screen->get_rgba_colormap;
$window->set_colormap($rgba_colormap) if $rgba_colormap;

my $drawing_area = Gtk2::DrawingArea->new;
$drawing_area->{_pixbuf_} = $pixbuf;
$window->add($drawing_area);

$drawing_area->signal_connect('expose-event' => \&on_draw);
$window->signal_connect(destroy => sub { Gtk2->main_quit; });

my $drag_info = {};
$window->add_events(['button-press-mask', 'button1-motion-mask']);
$window->signal_connect('button-press-event', sub {
    my ($widget, $event) = @_;
    if ($event->button == 1) {
        my ($root_x, $root_y) = $event->get_root_coords;
        my ($win_x, $win_y) = $widget->get_position;
        $drag_info->{x} = $root_x - $win_x;
        $drag_info->{y} = $root_y - $win_y;
    }
});
$window->signal_connect('motion-notify-event' => sub {
    my ($widget, $event) = @_;
    if (exists $drag_info->{x}) {
        my ($root_x, $root_y) = $event->get_root_coords;
        $widget->move($root_x - $drag_info->{x}, $root_y - $drag_info->{y});
    }
});
$window->signal_connect('button-release-event' => sub { $drag_info = {}; });

Glib::Timeout->add(1000, \&update_connections_and_redraw, $drawing_area);

$window->show_all;
update_connections_and_redraw($drawing_area);

Gtk2->main;

exit 0;