#!/usr/bin/perl

use strict;
use warnings;
use Gtk2 -init;
use Cairo;
use Glib;
use POSIX qw(tan log sqrt);
use constant PI => 3.14159265;
use Fcntl qw(SEEK_SET O_NONBLOCK F_GETFL F_SETFL); # For non-blocking sockets
use Net::IP;
use File::Basename;
use Text::CSV;
use Storable qw(retrieve store thaw);
use IO::Socket::INET;
use Errno qw(EAGAIN); # To check for non-blocking errors

# --- Configuration ---
my $listen_port = 6789;
my $home_dir = $ENV{'HOME'} or die "HOME environment variable not set";
my $resource_dir = "resources";
my $config_file = "connmapperlrc";
my $history_file = ".connmap-history.dat";
my $image_path = "$resource_dir/w1000-old.png";
my $ipv4_db_path = "$resource_dir/ipv4.csv";
my $geolite_file = "$resource_dir/GeoLite2-ASN-Blocks-IPv4.csv";
my $show_close_button = 1;
my $debug = 0;

# --- Global State ---
my $state = {
    current_points     => [],
    client_connections => {},
    history_points     => {},
    show_history       => 0,
    map_width          => 1000,
    map_height         => 500,
    dot_diameter       => 8,
    default_dot_color  => [1, 0.2, 0.2, 0.8],
    text_color         => [1, 1, 1, 0.9],
    ip2asn             => {},
    port_colors        => [],
    show_text_globally => 1,
    right_click_circle => undef,
    color_key_areas    => [],
    status_message     => "",
    status_message_timer_id => undef,
    show_help_text     => 1,
    help_text_area     => undef,
};

my $window;
my $drawing_area;
my @client_sockets;
my %client_buffers;
my $ip2asn_lookup_enabled = 0;
my $ip2asn_script_path = "";


#======================================================================
# Status Message Handling
#======================================================================
sub set_status_message {
    my ($message) = @_;
    print "$message\n";
    $state->{status_message} = $message;
    if (defined $state->{status_message_timer_id}) {
        Glib::Source->remove($state->{status_message_timer_id});
        $state->{status_message_timer_id} = undef;
    }
    $state->{status_message_timer_id} = Glib::Timeout->add(2000, sub {
        $state->{status_message} = "";
        $state->{status_message_timer_id} = undef;
        $drawing_area->queue_draw() if $drawing_area;
        return 0;
    });
    $drawing_area->queue_draw() if $drawing_area;
}

#======================================================================
# Configuration & History I/O
#======================================================================
sub save_history_points {
    eval { store($state->{history_points}, $history_file); };
    if ($@) {
        set_status_message("Error saving history: $@");
        return;
    }
    set_status_message("History saved to $history_file");
}

sub load_history_points {
    unless (-e $history_file) {
        set_status_message("History file not found: $history_file");
        return;
    }
    my $loaded_history;
    eval { $loaded_history = retrieve($history_file); };
    if ($@) {
        set_status_message("Error loading history: $@");
        return;
    }
    unless (ref($loaded_history) eq 'HASH') {
        set_status_message("Error: History file is corrupt or not valid.");
        return;
    }
    $state->{history_points} = $loaded_history;
    set_status_message("History loaded from $history_file");
    $drawing_area->queue_draw();
}

sub load_port_colors_from_config {
    my ($file) = @_;
    unless (-e $file) {
        warn "Configuration file not found: $file. Using default colors only.\n";
        return;
    }
    my $csv = Text::CSV->new({ binary => 1, auto_diag => 1, allow_loose_quotes => 1 });
    open(my $fh, "<:encoding(utf8)", $file) or do {
        warn "Could not open config file '$file': $!. Using defaults.\n";
        return;
    };
    my @loaded_colors;
    while (my $row = $csv->getline($fh)) {
        next if $row->[0] =~ /^\s*#/;
        next unless @$row == 8;
        my ($label, $type, $port, $r, $g, $b, $a, $show_text) = @$row;
        unless (defined $label && ($type eq 'local' || $type eq 'peer') && $port =~ /^\d+$/ && defined $show_text) {
            warn "Skipping invalid line in config file: " . join(',', @$row) . "\n";
            next;
        }
        push @loaded_colors, {
            label => $label, type => $type, port => int($port),
            color => [$r, $g, $b, $a], show_text => int($show_text),
        };
    }
    close $fh;
    $state->{port_colors} = \@loaded_colors;
    print "Loaded " . scalar(@loaded_colors) . " port color rules from $config_file.\n";
}

#======================================================================
# IP Geolocation and Coordinate Functions
#======================================================================
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];
}

sub find_location_on_disk {
    my ($fh, $ip_decimal, $file_size) = @_;
    my ($low_pos, $high_pos) = (0, $file_size);
    for (1..30) {
        last if ($high_pos - $low_pos) < 2;
        my $mid_pos = int($low_pos + ($high_pos - $low_pos) / 2);
        seek($fh, $mid_pos, 0); <$fh>;
        my $line = <$fh>;
        return undef unless defined $line;
        my ($start, $end, $lat, $lon) = split ',', $line;
        next unless defined $start and defined $end;
        if ($ip_decimal >= $start && $ip_decimal <= $end) {
            return { lat => $lat, lon => $lon };
        }
        $ip_decimal < $start ? ($high_pos = $mid_pos) : ($low_pos = $mid_pos);
    }
    return undef;
}

sub latlon_to_xy {
    my ($latitude, $longitude) = @_;
    my ($width, $height) = ($state->{map_width}, $state->{map_height});
    my ($xOffset, $yOffset) = (-29.0, 76.0);
    my $x = ($longitude + 180) * ($width / 360) + $xOffset;
    $x -= (($longitude/180) * 15) if $longitude > 90;
    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);
}

sub find_ip_in_geolite_asn {
    my ($geolite_csv_file, $ip_to_check) = @_;
    my $input_ip = Net::IP->new($ip_to_check);
    return unless $input_ip && $input_ip->version() == 4;
    my $input_ip_num = $input_ip->intip();
    open(my $fh, '<', $geolite_csv_file) or return;
    my $file_size = -s $fh;
    my ($low, $high) = (0, $file_size);
    my $found_line;
    $low = length(scalar <$fh>); # Skip header
    while ($low <= $high) {
        my $mid = int(($low + $high) / 2);
        seek($fh, $mid, SEEK_SET); <$fh> if $mid > 0;
        my $line = <$fh>;
        next unless defined $line && $line =~ /\S/;
        my ($network_cidr) = split /,/, $line, 2;
        my $network_ip = Net::IP->new($network_cidr);
        unless ($network_ip) { $high = $mid -1; next; }
        my $start_ip_num = $network_ip->intip();
        my $end_ip_num   = $start_ip_num + ($network_ip->size() - 1);
        if ($input_ip_num >= $start_ip_num && $input_ip_num <= $end_ip_num) {
            $found_line = $line; last;
        } elsif ($input_ip_num < $start_ip_num) {
            $high = $mid - 1;
        } else { $low = $mid + 1; }
    }
    close $fh;
    return unless defined $found_line;
    my @fields = split /,/, $found_line;
    return { network => $fields[0], asn => $fields[1], organization => $fields[2] || "N/A" };
}

sub get_point_label_text {
    my ($point) = @_;
    my $programname = $point->{programname} // "";
    if ($state->{ip2asn}->{$point->{ip}}) {
        return "$programname:$point->{local_port} -> $point->{ip}:$point->{port} - $state->{ip2asn}->{$point->{ip}}";
    } else {
        my $asn_hash = find_ip_in_geolite_asn($geolite_file, $point->{ip});
        if ($asn_hash) {
            my $asn_org_str = "$asn_hash->{asn} - $asn_hash->{organization}";
            $state->{ip2asn}->{$point->{ip}} = $asn_org_str;
            return "$programname:$point->{local_port} -> $point->{ip}:$point->{port} - $asn_org_str";
        } else {
            return "$programname:$point->{local_port} -> $point->{ip}:$point->{port}";
        }
    }
}

#======================================================================
# GTK, Drawing, and Timer Functions
#======================================================================
sub process_and_draw {
    my ($drawing_area) = @_;
    open my $db_fh, '<', $ipv4_db_path or return 1;
    my $db_size = -s $db_fh;
    
    my @connections;
    for my $fileno (keys %{$state->{client_connections}}) {
        push @connections, @{ $state->{client_connections}->{$fileno} || [] };
    }
    
    my @new_points;
    for my $conn (@connections) {
        my $ip_decimal = ip_to_decimal($conn->{ip});
        next unless $ip_decimal;
        my $location = find_location_on_disk($db_fh, $ip_decimal, $db_size);
        if ($location) {
            my $key = "$location->{lat},$location->{lon}";
            
            unless (exists $state->{history_points}->{$key}) {
                $state->{history_points}->{$key} = { count => 0, show_text_seen => 0 };
            }
            $state->{history_points}->{$key}->{count}++;

            if (!$state->{history_points}->{$key}->{show_text_seen}) {
                my $is_other_connection = 1;

                for my $rule (@{$state->{port_colors}}) {
                    if ( ($rule->{type} eq 'local' && $conn->{local_port} == $rule->{port}) ||
                         ($rule->{type} eq 'peer'  && $conn->{port} == $rule->{port}) )
                    {
                        $is_other_connection = 0;
                        if ($rule->{show_text}) {
                            $state->{history_points}->{$key}->{show_text_seen} = 1;
                        }
                        last;
                    }
                }

                if ($is_other_connection) {
                    $state->{history_points}->{$key}->{show_text_seen} = 1;
                }
            }
            
            push @new_points, { 
                ip => $conn->{ip}, port => $conn->{port}, 
                lat => $location->{lat}, lon => $location->{lon},
                local_ip => $conn->{local_ip}, local_port => $conn->{local_port},
                programname => $conn->{programname} 
            };
        }
    }
    close $db_fh;
    $state->{current_points} = \@new_points;
    
    $drawing_area->queue_draw();
    return 1; # Keep timer running
}


sub draw_color_key {
    my ($cr) = @_;
    my @key_items;
    my %seen_labels;

    for my $rule (@{$state->{port_colors}}) {
        my $label = $rule->{label};
        unless ($seen_labels{$label}) {
            push @key_items, { label => $label, color => $rule->{color} };
            $seen_labels{$label} = 1;
        }
    }
    push @key_items, { label => "Other", color => $state->{default_dot_color} };

    my $x_start = 15;
    my $line_height = 20;
    my $y_start = $state->{map_height} - (scalar(@key_items) * $line_height) - 10;
    my $box_size = 10;

    $cr->select_font_face('Sans', 'normal', 'bold');
    $cr->set_font_size(10);

    my $current_y = $y_start;
    my @key_areas;

    for my $item (@key_items) {
        push @key_areas, {
            x => $x_start, y => $current_y, w => $box_size, h => $box_size,
            label => $item->{label},
        };

        $cr->set_source_rgba(@{$item->{color}});
        $cr->rectangle($x_start, $current_y, $box_size, $box_size);
        $cr->fill;

        $cr->set_source_rgba(@{$state->{text_color}});
        $cr->move_to($x_start + 18, $current_y + 10);
        $cr->show_text($item->{label});

        $current_y += $line_height;
    }
    $state->{color_key_areas} = \@key_areas;
}

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;

    if ($state->{show_history}) {
        my $base_history_radius = $state->{dot_diameter} / 3;
        my $max_history_radius = $state->{dot_diameter} * 1.5;
        my $history_color = [1.0, 1.0, 1.0, 0.6];
        
        $cr->set_source_rgba(@{$history_color});

        for my $key (keys %{$state->{history_points}}) {
            my $history_entry = $state->{history_points}->{$key};
            next unless $history_entry->{show_text_seen};
            my ($lat, $lon) = split ',', $key;
            next unless defined $lat && defined $lon;
            my $count = $history_entry->{count};
            my $scaled_radius = $base_history_radius + log($count);
            $scaled_radius = $max_history_radius if $scaled_radius > $max_history_radius;
            my ($x, $y) = latlon_to_xy($lat, $lon);
            $cr->arc($x, $y, $scaled_radius, 0, 2 * PI);
            $cr->fill;
        }
    }

    my $radius = $state->{dot_diameter} / 2;
    for my $point (@{$state->{current_points}}) {
        my ($x, $y) = latlon_to_xy($point->{lat}, $point->{lon});
        my ($dot_color, $show_label) = ($state->{default_dot_color}, 1);
        for my $rule (@{$state->{port_colors}}) {
            if (($rule->{type} eq 'local' && $point->{local_port} == $rule->{port}) ||
                ($rule->{type} eq 'peer' && $point->{port} == $rule->{port})) {
                $dot_color = $rule->{color};
                $show_label = $rule->{show_text};
                last;
            }
        }
        $cr->set_source_rgba(@{$dot_color});
        $cr->arc($x, $y, $radius, 0, 2 * PI); $cr->fill;
        if ($show_label && $state->{show_text_globally}) {
            $cr->set_source_rgba(@{$state->{text_color}});
            $cr->select_font_face('Sans', 'normal', 'normal');
            $cr->set_font_size(8);
            $cr->move_to($x + $radius + 2, $y + $radius / 2);
            my $label = get_point_label_text($point);
            $cr->show_text($label);
        }
    }
    draw_color_key($cr);

    if (defined $state->{right_click_circle}) {
        my $circle = $state->{right_click_circle};
        $cr->new_path; 
        $cr->set_source_rgba(1, 1, 1, 0.9);
        $cr->set_line_width(1.5);
        $cr->arc($circle->{x}, $circle->{y}, $circle->{radius}, 0, 2 * PI);
        $cr->stroke;
    }

    if (defined $state->{status_message} && $state->{status_message} ne "") {
        my $message = $state->{status_message};
        $cr->select_font_face('Sans', 'normal', 'bold');
        $cr->set_font_size(14);
        my $extents = $cr->text_extents($message);
        my $padding = 10;
        my $rect_w = $extents->{width} + (2 * $padding);
        my $rect_h = $extents->{height} + (2 * $padding);
        my $rect_x = ($state->{map_width} - $rect_w) / 2;
        my $rect_y = $state->{map_height} - $rect_h - 5;
        $cr->set_source_rgba(0, 0, 0, 0.7);
        $cr->rectangle($rect_x, $rect_y, $rect_w, $rect_h);
        $cr->fill;
        my $text_x = $rect_x + $padding;
        my $text_y = $rect_y + $padding + $extents->{y_bearing} * -1;
        $cr->move_to($text_x, $text_y);
        $cr->set_source_rgba(@{$state->{text_color}});
        $cr->show_text($message);
    }

    if ($state->{show_help_text}) {
        my $help_text = "c: reload config | " . "press: t, h, s, l, esc, or right click [click here to close help]";
        $cr->select_font_face('Sans', 'normal', 'normal');
        $cr->set_font_size(9);
        my $help_extents = $cr->text_extents($help_text);
        my $help_padding = 5;
        my $help_x = $state->{map_width} - $help_extents->{width} - $help_padding;
        my $help_y = $state->{map_height} - $help_padding;
        
        $state->{help_text_area} = {
            x => $help_x,
            y => $help_y - $help_extents->{height},
            w => $help_extents->{width},
            h => $help_extents->{height},
        };

        $cr->set_source_rgba(@{$state->{text_color}});
        $cr->move_to($help_x, $help_y);
        $cr->show_text($help_text);
    }

    return 1;
}

sub create_info_popup {
    my ($info_lines_ref) = @_;
    my $popup = Gtk2::Window->new('toplevel');
    $popup->set_title("Nearby Connections");
    $popup->set_default_size(450, 250);
    $popup->set_position('center-on-parent');
    $popup->set_transient_for($window);
    $popup->set_destroy_with_parent(1);
    $popup->set_border_width(5);

    my $main_vbox = Gtk2::VBox->new(0, 5);
    $popup->add($main_vbox);

    if ($ip2asn_lookup_enabled) {
        my $lookup_hbox = Gtk2::HBox->new(0, 5);
        my $ip_label = Gtk2::Label->new("IP:");
        my $ip_entry = Gtk2::Entry->new();
        my $lookup_button = Gtk2::Button->new("Look-up");

        $lookup_hbox->pack_start($ip_label, 0, 0, 5);
        $lookup_hbox->pack_start($ip_entry, 1, 1, 0);
        $lookup_hbox->pack_start($lookup_button, 0, 0, 5);

        $main_vbox->pack_start($lookup_hbox, 0, 0, 5);

        $lookup_button->signal_connect(clicked => sub {
            my $ip = $ip_entry->get_text();
            $ip =~ s/^\s+|\s+$//g;

            unless ($ip =~ /^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$/ || Net::IP->new($ip)) {
                warn "Invalid IP address format entered: '$ip'\n";
                return;
            }
            
            my $command = "perl $ip2asn_script_path \"$ip\"";
            my $output = qx($command);

            my $output_popup = Gtk2::Window->new('toplevel');
            $output_popup->set_title("Lookup: $ip");
            $output_popup->set_default_size(400, 300);
            $output_popup->set_position('center-on-parent');
            $output_popup->set_transient_for($popup);
            $output_popup->set_destroy_with_parent(1);
            
            my $output_sw = Gtk2::ScrolledWindow->new(undef, undef);
            $output_sw->set_policy('automatic', 'automatic');
            
            my $output_tv = Gtk2::TextView->new();
            $output_tv->set_editable(0);
            $output_tv->get_buffer->set_text($output || "No output from script for $ip.");
            
            $output_sw->add($output_tv);
            $output_popup->add($output_sw);
            $output_popup->show_all;
        });
    }

    my $scrolled_window = Gtk2::ScrolledWindow->new(undef, undef);
    $scrolled_window->set_policy('automatic', 'automatic');
    $scrolled_window->set_shadow_type('in');
    
    my $textview = Gtk2::TextView->new;
    $textview->set_editable(0);
    $textview->set_cursor_visible(0);
    $textview->set_wrap_mode('word');
    my $buffer = $textview->get_buffer;
    my $text = join("\n", @$info_lines_ref);
    $buffer->set_text($text);

    $scrolled_window->add($textview);
    $main_vbox->pack_start($scrolled_window, 1, 1, 0);
    
    $popup->show_all;
}

sub handle_right_click {
    my ($click_x, $click_y) = @_;
    my @info_strings;
    my $radius = 15;
    my $radius_squared = $radius * $radius;   
    for my $point (@{$state->{current_points}}) {
        my ($px, $py) = latlon_to_xy($point->{lat}, $point->{lon});
        my $dist_squared = ($px - $click_x)**2 + ($py - $click_y)**2;
        if ($dist_squared <= $radius_squared) {
            if ($state->{show_text_globally}) {
                my $full_label = get_point_label_text($point);
                push @info_strings, "$full_label GPS: $point->{lat}, $point->{lon}";
            } else {
                push @info_strings, "IP: $point->{ip} GPS: $point->{lat}, $point->{lon}";
            }
        }
    }
    
    create_info_popup(\@info_strings) if @info_strings;
    
    $state->{right_click_circle} = { x => $click_x, y => $click_y, radius => $radius };
    $drawing_area->queue_draw();

    Glib::Timeout->add(500, sub {
        $state->{right_click_circle} = undef;
        $drawing_area->queue_draw();
        return 0;
    });
}

#======================================================================
# Network Handling Subroutines (Glib::IO-free)
#======================================================================

sub set_non_blocking {
    my ($socket) = @_;
    my $flags = fcntl($socket, F_GETFL, 0) or die "Can't get flags for socket: $!";
    fcntl($socket, F_SETFL, $flags | O_NONBLOCK) or die "Can't set non-blocking flag: $!";
}

sub handle_network_events {
    # 1. Check for new connections on the server socket
    if (my $client_sock = $main::server_sock->accept) {
        my $client_ip = $client_sock->peerhost();
        set_non_blocking($client_sock);
        push @client_sockets, $client_sock;
        my $fileno = fileno($client_sock);
        $client_buffers{$fileno} = { buf => '', len => -1 };
        set_status_message("Accepted new client from $client_ip (fd:$fileno)");
    }

    # 2. Check all existing clients for data
    my @sockets_to_keep;
    foreach my $sock (@client_sockets) {
        my $fileno = fileno($sock);
        my $buf;
        my $bytes_read = sysread($sock, $buf, 16384);

        if (defined $bytes_read) {
            if ($bytes_read > 0) {
                $client_buffers{$fileno}->{buf} .= $buf;
                
                while (1) {
                    my $b = $client_buffers{$fileno};
                    if ($b->{len} == -1) {
                        if (length($b->{buf}) >= 4) {
                            $b->{len} = unpack('N', substr($b->{buf}, 0, 4, ''));
                        } else { last; }
                    }
                    
                    if ($b->{len} != -1 && length($b->{buf}) >= $b->{len}) {
                        my $payload = substr($b->{buf}, 0, $b->{len}, '');
                        my $connections_ref = eval { thaw($payload) };
                        if ($@) {
                            warn "Failed to deserialize data from client (fd:$fileno): $@\n";
                        } else {
                            $state->{client_connections}->{$fileno} = $connections_ref;
                        }
                        $b->{len} = -1; 
                    } else { last; }
                }
                push @sockets_to_keep, $sock;
            } else {
                set_status_message("Client disconnected (fd:$fileno)");
                close $sock;
                delete $state->{client_connections}->{$fileno};
                delete $client_buffers{$fileno};
            }
        } elsif ($! != EAGAIN) {
            set_status_message("Client connection error (fd:$fileno): $!");
            close $sock;
            delete $state->{client_connections}->{$fileno};
            delete $client_buffers{$fileno};
        } else {
            push @sockets_to_keep, $sock;
        }
    }
    
    @client_sockets = @sockets_to_keep;
    return 1;
}

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

$ip2asn_script_path = "ip2asn.pl";
if (-e $ip2asn_script_path && -x $ip2asn_script_path) {
    $ip2asn_lookup_enabled = 1; 
    print "Found executable '$ip2asn_script_path', lookup feature is enabled.\n";
} else { 
    print "Did not find executable 'ip2asn.pl', lookup feature is disabled.\n"; 
}

load_port_colors_from_config($config_file);

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 $@;

$window = Gtk2::Window->new('toplevel');
$window->set_title("connmap-gtk2-server");
$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;
if ($rgba_colormap) { $window->set_colormap($rgba_colormap); }
else { print "DEBUG: No RGBA colormap available. Transparency may not work.\n"; }

my $fixed_container = Gtk2::Fixed->new;
$drawing_area = Gtk2::DrawingArea->new;
$drawing_area->{_pixbuf_} = $pixbuf;
$drawing_area->set_size_request($state->{map_width}, $state->{map_height});
$fixed_container->put($drawing_area, 0, 0);

if ($show_close_button) {
    my $close_button = Gtk2::Button->new_from_stock('gtk-close');
    $close_button->set_relief('none');
    my $req = $close_button->size_request;
    $fixed_container->put($close_button, $state->{map_width} - $req->width - 5, 5);
    $close_button->signal_connect(clicked => sub { Gtk2->main_quit; });
}

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

# --- THIS BLOCK IS MODIFIED ---
$window->signal_connect('key-press-event' => sub {
    my ($widget, $event) = @_;
    my $keyname = Gtk2::Gdk->keyval_name($event->keyval);
    if ($keyname eq 'h') { 
        $state->{show_history} = !$state->{show_history}; 
        set_status_message("History view: " . ($state->{show_history} ? "ON" : "OFF")); 
    } elsif ($keyname eq 't') { 
        $state->{show_text_globally} = !$state->{show_text_globally}; 
        set_status_message("Global text labels: " . ($state->{show_text_globally} ? "ON" : "OFF")); 
    } elsif ($keyname eq 's') { 
        save_history_points(); 
    } elsif ($keyname eq 'l') { 
        load_history_points(); 
    } elsif ($keyname eq 'c') { # New key handler
        load_port_colors_from_config($config_file);
        set_status_message("Reloaded config from $config_file");
    } elsif ($keyname eq 'Escape') { 
        set_status_message("Exiting..."); 
        Gtk2->main_quit; 
    }
});

my $drag_info = {};
$window->add_events(['button-press-mask', 'button1-motion-mask', 'button-release-mask', 'key-press-mask']);

$window->signal_connect('button-press-event' => sub {
    my ($widget, $event) = @_;
    my ($click_x, $click_y) = $event->get_coords;

    if ($event->button == 1) { # Left click

        for my $area (@{$state->{color_key_areas}}) {
            if ($click_x >= $area->{x} && $click_x <= ($area->{x} + $area->{w}) &&
                $click_y >= $area->{y} && $click_y <= ($area->{y} + $area->{h}))
            {
                my $label_clicked = $area->{label};
                if ($label_clicked eq 'Other') {
                    set_status_message("'Other' is not a toggleable rule");
                    return 1;
                }
                my $new_state;
                for my $rule (@{$state->{port_colors}}) {
                    if ($rule->{label} eq $label_clicked) {
                        $rule->{show_text} = !$rule->{show_text};
                        $new_state = $rule->{show_text};
                    }
                }
                set_status_message(sprintf("Rule '%s' text: %s", $label_clicked, ($new_state ? "ON" : "OFF")));
                $drawing_area->queue_draw();
                return 1;
            }
        }

        if ($state->{show_help_text} && defined $state->{help_text_area}) {
            my $area = $state->{help_text_area};
            if ($click_x >= $area->{x} && $click_x <= ($area->{x} + $area->{w}) &&
                $click_y >= $area->{y} && $click_y <= ($area->{y} + $area->{h}))
            {
                $state->{show_help_text} = 0;
                $drawing_area->queue_draw();
                return 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;
        return 1;
    }
    elsif ($event->button == 3) { # Right click
        handle_right_click($click_x, $click_y);
        return 1;
    }
    return 0;
});

$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 = {};
});

# Set up the listening socket and make it non-blocking
our $server_sock = IO::Socket::INET->new(
    LocalPort => $listen_port,
    Type      => SOCK_STREAM,
    Reuse     => 1,
    Listen    => 10,
) or die "Could not create server socket on port $listen_port: $!\n";
set_non_blocking($server_sock);

# A short timeout for responsive network I/O
Glib::Timeout->add(100, \&handle_network_events);

# A longer timeout for less frequent drawing updates
Glib::Timeout->add(1000, \&process_and_draw, $drawing_area);

print "Server listening on port $listen_port...\n";
$window->show_all;
process_and_draw($drawing_area);
Gtk2->main;

# clean up sockets on exit
close($server_sock);
foreach my $sock (@client_sockets) {
    close($sock);
}

exit 0;