#!/usr/bin/perl
############################################################################
# Copyright 2006 University of Washington. All Rights Reserved.
# Licensing and technical questions: tv-streaming@cac.washington.edu
# Created by Matt Hodge, SMT, Computing & Communications.
############################################################################

# TODO: d type

use strict;
$| = 1;

use Socket;
use Fcntl;
use Errno;

if(! @ARGV) {
    print "Usage:

    mpe-cmd.pl cmd host:port <cmd> ...
    mpe-cmd.pl stats host:port <wkr_id> ... [ -- <stat_id> ... ]

";
    exit 0;
}

my $retval;
my $mode = shift(@ARGV);
if($mode eq 'cmd') { $retval = main_cmd(@ARGV); }
elsif($mode eq 'stats') { $retval = main_stats(@ARGV); }
else { die "mode invalid"; }
exit($retval ? 0 : 1);

sub main_cmd {
    my ($addr, @cmds) = @_;
    
    @cmds || die "wrong number args";
    
    my $sock = socket_connect($addr);

    my $status = 'CONTINUE';
    my %mpe_state;
    my @rest;
    mpe_parse_init(\%mpe_state, 1);

    for my $cmd (@cmds) {
	print "SEND $cmd\n";
	
	($status, @rest) = send_recv($sock, "$cmd\n", $status, \%mpe_state);
	if($status eq 'OK') {
	    print "RECV ".mpe_string($rest[0])."\n";
	} else { die join(' ', $status, @rest); }
    }
    close($sock);
}

sub main_stats {
    my ($addr, $interval, @args) = @_;

    @args || die "wrong number args";

    my @wkr_ids;
    my @stat_specs;
    for(my $i = 0; ; $i++) {
	if($i == @args) {
	    @wkr_ids = @args;
	    last;
	} elsif($args[$i] eq '--') {
	    @stat_specs = splice(@args, $i + 1);
	    @wkr_ids = splice(@args, 0, $i);
	    last;
	}
    }

    my $sock = socket_connect($addr);

    my $status = 'CONTINUE';
    my %mpe_state;
    my @rest;
    mpe_parse_init(\%mpe_state, 1);

    my $cmd = '(stats_init '.join(' ', @wkr_ids).")\n";
    ($status, @rest) = send_recv($sock, $cmd, $status, \%mpe_state);
    if($status eq 'OK') {
	my %stats0;
	my %stats1;
	for(;;) {
	    ($status, @rest) = send_recv($sock, "(stats_get)\n", $status, \%mpe_state);
	    if($status eq 'OK') {
		if(ref($rest[0])) {
		    if($rest[0]->[0] eq 'ok') {
			my (undef, @mpes) = @{$rest[0]};
			stats_flatten(\%stats0, '', @mpes);
			stats_output(\%stats0, \%stats1, @stat_specs);
			%stats1 = %stats0;
		    } else { die "STATS_GET ERROR ".mpe_string($rest[0]); }
		} else { die "STATS_GET ERROR ".mpe_string($rest[0]); }
	    } else { die "STATS_GET ERROR ".mpe_string($rest[0]); }
	    sleep($interval);
	}
    } else { die "STATS_INIT ERROR ".mpe_string($rest[0]); }
    close($sock);
}

###################################################################
# stats stuff
###################################################################

sub stats_flatten {
    my ($stats, $prefix, @mpes) = @_;
    
    for my $mpe (@mpes) {
	if(ref($mpe)) {
	    if($mpe->[0] eq 'list') {
		my (undef, $id, @mpes) = @$mpe;
		stats_flatten($stats, $prefix.$id.'.', @mpes);
	    } elsif($mpe->[0] =~ m/^(u64|u32|i32|d)$/s) {
		$stats->{$prefix.$mpe->[1]} = $mpe->[2];
	    } else { die "$mpe->[0]: type invalid"; }
	} else { die "$mpe: expected list"; }
    }
}

sub stats_output {
    my ($stats0, $stats1, @specs) = @_;
    
    my @keys0 = (sort { $a cmp $b } (keys(%$stats0)));
    my @keys1 = (sort { $a cmp $b } (keys(%$stats1)));

    my @ids;
    if(! @specs) {
	@ids = @keys0;
    } else {
	for my $id (@keys0) {
	    if(grep { $id =~ qr/$_/ } @specs) { push(@ids, $id) };
	}
    }
	
	    
    my $time = $stats0->{'time'};
    if(join(' ', @keys0) eq join(' ', @keys1)) {
	my $time_diff = $stats0->{'time'} - $stats1->{'time'};
	for my $id (@ids) {
	    my $diff = $stats0->{$id} - $stats1->{$id};
	    print $id.' = '.$stats0->{$id}.' ['.($stats0->{$id} / $time).'] '.
		$diff.' ['.($diff / $time_diff)."]\n";
	}
    } else {
	for my $id (@ids) {
	    print $id.' = '.$stats0->{$id}.' ['.($stats0->{$id} / $time)."]\n";
	}
    }
}

###################################################################
# composite functions
###################################################################

sub send_recv {
    my ($sock, $send_buf, $status, $mpe_state) = @_;
    
    my $z = length($send_buf);
    if(! defined(my $x = send($sock, $send_buf, 0))) {
	return ('SEND_ERROR', $!);
    } elsif($x != $z) {
	return ('SEND_ERROR', "$x/$z");
    }
    
    return(recv_mpe($sock, $status, $mpe_state));
}

sub recv_mpe {
    my ($sock, $status, $mpe_state) = @_;
    my @rest;
    
    if($status ne 'CONTINUE') { goto loop_post_recv; }

  loop:
    my $buf;
    if(! defined(recv($sock, $buf, 8192, 0))) {
	return ('RECV_ERROR', $!);
    } elsif(! length($buf)) {
	mpe_parse_add($mpe_state, '', 1);
    } else {
	mpe_parse_add($mpe_state, $buf);
    }

  loop_post_recv:
    ($status, @rest) = mpe_parse($mpe_state);
    if($status eq 'CONTINUE') { goto loop; }
    
    return ($status, @rest);
}

###################################################################
# mpe_parse
###################################################################

sub mpe_parse_init {
    my ($state, $single) = @_;

    $state->{'single'} = $single;
    $state->{'src'} = '';
    $state->{'src_end'} = undef;
    $state->{'src_pos'} = 0;
    $state->{'tok_state'} = 'WS';
    # $state->{'tok_start'}
    # $state->{'tok'}
    $state->{'stack'} = [];
    $state->{'syms'} = [];
}

sub mpe_parse_add {
    my ($state, $part, $end) = @_;
    $state->{'src'} .= $part;
    $state->{'src_end'} = $end;
}

sub mpe_parse {
    my ($state) = @_;
    
    for(;;) {
	for( ; ; $state->{'src_pos'}++) {
	    if($state->{'src_pos'} >= length($state->{'src'})) {
		if(! $state->{'src_end'}) { return ('CONTINUE'); }
		elsif($state->{'tok_state'} eq 'NORMAL') { goto sym_str; }
		elsif($state->{'tok_state'} eq 'WS') {
		    if(@{$state->{'stack'}}) { return ('UNMATCHED_LP'); }
		    elsif(@{$state->{'syms'}}) { goto ok; }
		    else { return ('DONE'); }
		} else { return ('TOK_INCOMPLETE'); }
		
	    }
	    my $c = substr($state->{'src'}, $state->{'src_pos'}, 1);
	    if($state->{'tok_state'} eq 'WS') {
		if($c !~ m/\s/) {
		    $state->{'tok_state'} = 'NORMAL';
		    $state->{'tok_start'} = $state->{'src_pos'};
		    $state->{'tok'} = '';
		    goto tok_normal;
		}
	    } elsif($state->{'tok_state'} eq 'NORMAL') {
	      tok_normal:
		if($c eq '"') { $state->{'tok_state'} = 'DQ'; }
		elsif($c eq "'") { $state->{'tok_state'} = 'SQ'; }
		elsif($c eq "\\") { $state->{'tok_state'} = 'BS'; }
		elsif($c eq '(') {
		    if($state->{'src_pos'} > $state->{'tok_start'}) { goto sym_str; }
		    $state->{'src_pos'}++;
		    goto sym_lp;
		} elsif($c eq ')') {
		    if($state->{'src_pos'} > $state->{'tok_start'}) { goto sym_str; }
		    $state->{'src_pos'}++;
		    goto sym_rp;
		} elsif($c =~ m/\s/) {
		    goto sym_str;
		} else { $state->{'tok'} .= $c; }
	    } elsif($state->{'tok_state'} eq 'BS') {
		$state->{'tok'} .= $c;
		$state->{'tok_state'} = 'NORMAL';
	    } elsif($state->{'tok_state'} eq 'BS_DQ') {
		$state->{'tok'} .= $c;
		$state->{'tok_state'} = 'DQ';
	    } elsif($state->{'tok_state'} eq 'DQ') {
		if($c eq '"') { $state->{'tok_state'} = 'NORMAL'; }
		elsif($c eq "\\") { $state->{'tok_state'} = 'BS_DQ'; }
		else { $state->{'tok'} .= $c; }
	    } elsif($state->{'tok_state'} eq 'SQ') {
		if($c eq "'") { $state->{'tok_state'} = 'NORMAL'; }
		else { $state->{'tok'} .= $c; }
	    }
	}
	
      sym_str:
	if(@{$state->{'stack'}}) {
	    push(@{$state->{'stack'}->[$#{$state->{'stack'}}]}, $state->{'tok'});
	} else {
	    push(@{$state->{'syms'}}, $state->{'tok'});
	}
	goto sym_done;
	
      sym_lp:
	push(@{$state->{'stack'}}, []);
	goto sym_done;
	
      sym_rp:
	if(defined(my $sym = pop(@{$state->{'stack'}}))) {
	    if(@{$state->{'stack'}}) {
		push(@{$state->{'stack'}->[$#{$state->{'stack'}}]}, $sym);
	    } else {
		push(@{$state->{'syms'}}, $sym);
	    }
	} else { return ('UNMATCHED_RP'); }
	goto sym_done;

      sym_done:
	$state->{'tok_state'} = 'WS';
	if((! @{$state->{'stack'}}) && $state->{'single'}) { goto ok; }
    }
  ok:
    my $syms = $state->{'syms'};
    $state->{'syms'} = [];
    return ('OK', @$syms);
}    

############################################################################
# mpe_string
############################################################################

sub mpe_string {
    my ($mpe) = @_;
    
    if(ref($mpe)) {
	return '( '.join(' ', map { mpe_string($_) } @$mpe).' )';
    } elsif(! length($mpe)) {
	return '""';
    } elsif($mpe =~ m|[^A-Za-z0-9_./-]|s) {
	$mpe =~ s/[\"\\]/\\$&/sg;
	return '"'.$mpe.'"';
    } else {
	return $mpe;
    }
}

############################################################################
# socket functions
############################################################################

sub packaddr {
    my ($host, $port) = split(/:/, $_[0]);
    my (undef, undef, undef, undef, $hostaddr) = gethostbyname($host);
    return pack('S n a4 x8', &AF_INET, $port, $hostaddr);
}

sub socket_connect {
    my ($addr) = @_;
    
    local *SOCK;
    if(socket(SOCK, &PF_INET, &SOCK_STREAM, 0)) {
	if(connect(SOCK, packaddr($addr))) {
	    return *SOCK;
	} else { print STDERR "$!, connect failed\n"; }
	close(SOCK);
    } else { print STDERR "$!, socket failed\n"; }
    return;
}