# $Id: traffic.tcl 1746 2009-03-26 19:17:54Z sergei $
# This plugin counts in and out bytes per JID.
#
# Bugs and feature request send to feez@jabber.ru

package require msgcat

namespace eval traffic {
    ::msgcat::mcload [file join [file dirname [info script]] msgs]

    if {![::plugins::is_registered traffic]} {
	::plugins::register traffic \
			    -namespace [namespace current] \
			    -source [info script] \
			    -description [::msgcat::mc "Whether the Traffic Counter plugin is loaded."] \
			    -loadcommand [namespace code load] \
			    -unloadcommand [namespace code unload]
	return
    }

    # Options
    custom::defgroup Plugins [::msgcat::mc "Plugins options."] -group Tkabber

    custom::defgroup Traffic [::msgcat::mc "Traffic Counter plugin options."] -group Plugins

    custom::defvar options(show_nick) 1 \
        [::msgcat::mc "Show nick instead of JID."] \
        -type boolean -group Traffic \
	-command [namespace code clear_and_refresh]

    custom::defvar options(human_readable) 1 \
        [::msgcat::mc "Use human-readable counters with K, M, G."] \
        -type boolean -group Traffic
}

proc traffic::load {} {
    variable stats
    variable jids
    variable names
    variable columns
    variable sort_order
    variable current_xlib
    variable wished_xlib
    variable timer
    variable timer_state
    variable timer_id
    variable timer_names

    array set stats {}
    array set jids {}

    # for table
    array set names [list  \
			presence,in  	[::msgcat::mc "Pres IN"] \
			presence,out 	[::msgcat::mc "Pres OUT"] \
                        iq,in 		[::msgcat::mc "Iq IN"] \
                        iq,out 		[::msgcat::mc "Iq OUT"] \
                        message,in 	[::msgcat::mc "Msg IN"] \
                        message,out 	[::msgcat::mc "Msg OUT"] \
                        total,in 	[::msgcat::mc "All IN"] \
                        total,out 	[::msgcat::mc "All OUT"]]
    set columns [list presence,in presence,out iq,in iq,out message,in message,out total,in total,out]
    set sort_order total,in

    set current_xlib {}
    set wished_xlib [list {} None]

    set timer 5
    set timer_state 0
    set timer_id {}
    array set timer_names [list 0 [::msgcat::mc "On"] 1 [::msgcat::mc "Off"]]

    hook::add log_hook [namespace current]::handle_inout
    hook::add finload_hook [namespace current]::setup_menu
    hook::add connected_hook [namespace code [list update_connections on]]
    hook::add disconnected_hook [namespace code [list update_connections off]]
    hook::add save_session_hook [namespace current]::save_session

    setup_menu
}

proc traffic::unload {} {
    variable stats
    variable jids
    variable names
    variable columns
    variable sort_order
    variable current_xlib
    variable wished_xlib
    variable timer
    variable timer_state
    variable timer_id
    variable timer_names

    setup_menu

    if {[winfo exists .traffic_stats]} {
	destroy_win .traffic_stats
    }

    hook::remove log_hook [namespace current]::handle_inout
    hook::remove finload_hook [namespace current]::setup_menu
    hook::remove connected_hook [namespace code [list update_connections on]]
    hook::remove disconnected_hook [namespace code [list update_connections off]]
    hook::remove save_session_hook [namespace current]::save_session

    catch {unset stats}
    catch {unset jids}
    catch {unset names}
    catch {unset columns}
    catch {unset sort_order}
    catch {unset current_xlib}
    catch {unset wished_xlib}
    catch {unset timer}
    catch {unset timer_state}
    catch {unset timer_id}
    catch {unset timer_names}
}

# ############################
# Handlers

proc traffic::handle_inout {xlib dir type xmldata} {
    variable stat
    variable jids
    variable columns
    variable names

    switch -- $dir {
	output {
	    set prefix out
	    set attr to
	}
	default {
	    set prefix in
	    set attr from
	}
    }

    switch -- $type {
	xml {
	    set size [string bytelength [::xmpp::xml::toText $xmldata]]

	    ::xmpp::xml::split $xmldata tag xmlns attrs cdata subels

	    set jid [::xmpp::xml::getAttr $attrs $attr]
	    if {[cequal $jid ""]} {
		set jid SERVER
	    } else {
		catch { set jid [::xmpp::jid::stripResource $jid] }
	    }
	}
	default {
	    set tag ""
	    set size [string bytelength $xmldata]
	    set jid SERVER
	}
    }

    if {![info exists jids($xlib,$jid)]} {

	# create
        set jids($xlib,$jid) 1
	lappend jids($xlib,jids) $jid

	# initialize counters
	foreach column $columns {
	    set stat($xlib,$jid,$column) 0
	}
    }

    # increment counters
    if {[info exists names($tag,$prefix)]} {
        incr stat($xlib,$jid,$tag,$prefix) $size
    }
    incr stat($xlib,$jid,total,$prefix) $size
}

# ############################
# GUI

proc traffic::get_connections {} {
    set res [list [list {} None]]
    foreach con [connections] {
	lappend res [list $con [connection_jid $con]]
    }
    return $res
}

proc traffic::open_window {} {
    variable columns
    variable names
    variable current_xlib
    variable timer_state
    variable timer_names

    # create widgets
    set w .traffic_stats
    if {[winfo exists $w]} {
	raise_win $w
	return
    }

    # base widget
    add_win $w -title [::msgcat::mc "Traffic statistics"] \
	-tabtitle [::msgcat::mc "Traffic"] \
	-class Traffic \
	-raise 1

    # button box
    set tools [frame $w.tools -borderwidth 5]
    pack $tools -side top -anchor w -fill x

    button $tools.refresh -text [::msgcat::mc "Refresh"] \
        -command [list [namespace current]::refresh]
    pack $tools.refresh -side right -anchor w

    set connections [get_connections]
    ComboBox $tools.connection -textvariable [namespace current]::wished_xlib \
	-values $connections -height [expr [llength $connections] + 1] -editable 0 \
	-modifycmd [list [namespace current]::refresh]
    pack $tools.connection -side right -anchor w -expand yes -fill x
    unset connections

    button $tools.show_xml -text [::msgcat::mc "Show XML"] \
        -command [list ::plugins::rawxml::open_window]
    pack $tools.show_xml -side left -anchor w

    button $tools.save -text [::msgcat::mc "Save"] \
        -command [list [namespace current]::save_stat]
    pack $tools.save -side left -anchor w

    button $tools.reset -text [::msgcat::mc "Reset"] \
        -command [list [namespace current]::reset]
    pack $tools.reset -side left -anchor w


    pack [label $tools.timerlab -text [::msgcat::mc "Timer:"]] -side left -anchor w

    Spinbox $tools.timerspin 1 100000 1 [namespace current]::timer
    $tools.timerspin configure -width 4
    pack $tools.timerspin -side left -anchor w

    button $tools.timer -text $timer_names($timer_state) \
        -command [list [namespace current]::turn_timer]
    pack $tools.timer -side left -anchor w


    # create grid
    set sw [ScrolledWindow $w.sw -relief sunken -bd 1]
    pack $sw -side top -fill both -expand yes

    set sf [ScrollableFrame $w.sf]
    $sw setwidget $sf
    set f [$sf getframe]

    bindscroll $f $sf

    # fill grid header
    grid [label $f.titlejid -text [::msgcat::mc "JID"] -width 20] -row 0 -column 0 -sticky w
    bindscroll $f.titlejid $sf
    set i 1
    foreach col $columns {

        set b [Button $f.titlelabel$i -text $names($col) \
		    -command [list [namespace current]::sort $col]]
        grid $b -row 0 -column $i -sticky we
        bindscroll $f.titlelabel$i $sf
        incr i
    }

    refresh
}

proc traffic::refresh {} {
    variable stat
    variable jids
    variable columns
    variable wished_xlib
    variable current_xlib
    variable options

    set w .traffic_stats
    if {![winfo exists $w]} return
    set f [$w.sf getframe]


    # clear if connection was changed
    if {![cequal [lindex $wished_xlib 0] $current_xlib]} {
	clear_table $f
	set current_xlib [lindex $wished_xlib 0]
    }

    if {$current_xlib == {}} return

    # init sum counters
    array set sum {}
    foreach col $columns {
	set sum($col) 0
    }

    # fill the table
    set row 1
    foreach jid [lsort -decreasing -command compare $jids($current_xlib,jids)] {

	set tag [jid_to_tag $jid]

	# set first column - jid
	if {![winfo exists $f.jid$tag]} {
	    label $f.jid$tag -text $jid

	    if {$options(show_nick)} {
		set nick [roster::itemconfig $current_xlib [roster::find_jid $current_xlib $jid] -name]
		if {$nick != ""} {
		    $f.jid$tag configure -text $nick
		}
	    }
	}
	grid $f.jid$tag -row $row -column 0 -sticky w
       bindscroll $f.jid$tag $w.sf

        set i 1

	# set other columns - counters
	foreach col $columns  {

	    if {![winfo exists $f.value$i$tag]} {
		label $f.value$i$tag
	    }
	    set value $stat($current_xlib,$jid,$col)

	    incr sum($col) $value
	    $f.value$i$tag configure -text [convert_to_hr $value 10]
	    grid $f.value$i$tag -row $row -column $i -sticky e
           bindscroll $f.value$i$tag $w.sf

    	    incr i
	}

	incr row
    }

    # add sum
    if {![winfo exists $f.sumname]} {
	label $f.sumname -text [::msgcat::mc "Total"]
    }
    grid $f.sumname -row $row -column 0 -sticky w
    bindscroll $f.sumname $w.sf

    set i 1
    foreach col $columns {
	if {![winfo exists $f.sumval$i]} {
	    label $f.sumval$i
	}
	$f.sumval$i configure -text [convert_to_hr $sum($col) 10]
	grid $f.sumval$i -row $row -column $i -sticky e
	bindscroll $f.sumval$i $w.sf
	incr i
    }
}

proc traffic::convert_to_hr {value limit} {
    variable options

    if {!$options(human_readable)} {
	return $value
    }

    set cur ""
    foreach c {K M G} {
	if {($value / 1024) > $limit} {
	    set value [expr {$value / 1024}]
	    set cur $c
	} else {
	    break
	}
    }
    return $value$cur
}

proc traffic::clear_table {f} {

    set max [expr [string length $f.title] - 1]
    foreach slave [grid slaves $f] {
	if {![cequal [string range $slave 0 $max] $f.title]} {
	    grid remove $slave
	    destroy $slave
	}
    }
}

proc traffic::clear_and_refresh {args} {

    set w .traffic_stats
    if {![winfo exists $w]} return
    set f [$w.sf getframe]

    clear_table $f
    refresh
}

proc traffic::update_connections {type xlib} {
    variable wished_xlib

    set w .traffic_stats
    if {![winfo exists $w]} return

    set connections [get_connections]

    if {[cequal $type off] && [lsearch $connections $wished_xlib] < 0} {
	set wished_xlib [list {} None]
    }

    $w.tools.connection configure -values $connections \
				  -height [expr [llength $connections] + 1]

    if {[cequal $type on] && [cequal $wished_xlib [list {} None]]} {
	set wished_xlib [list $xlib [connection_jid $xlib]]
    }

    refresh
}

proc traffic::tick {} {
    variable timer_id
    variable timer

    refresh
    set timer_id [after [expr $timer * 1000] [list [namespace current]::tick]]
}

proc traffic::turn_timer {args} {
    variable timer_state
    variable timer_names
    variable timer_id
    variable timer

    set timer_state [expr 1 - $timer_state]

    if {$timer_state}  {
	set timer_id [after [expr $timer * 1000] [list [namespace current]::tick]]
    } else {
	after cancel $timer_id
    }

    set w .traffic_stats
    if {![winfo exists $w]} return

    $w.tools.timer configure -text $timer_names($timer_state)
}

# ############################
# Other functions

proc traffic::compare {arg1 arg2} {
    variable stat
    variable sort_order
    variable current_xlib

    return [expr $stat($current_xlib,$arg1,$sort_order) - \
		 $stat($current_xlib,$arg2,$sort_order)]
}

proc traffic::sort {tag} {
    variable sort_order
    set sort_order $tag

    refresh
}

proc traffic::save_stat {} {
    variable names
    variable columns
    variable stat
    variable current_xlib
    variable jids

    set filename [tk_getSaveFile \
                      -initialdir $::configdir]
    if {$filename != ""} {
	set fd [open $filename w]

	# Header
	set str {}
	lappend str "JID"
	foreach col $columns {
	    lappend str $names($col)
	}
	puts $fd [join $str "\t"]

	# Data
	foreach jid $jids($current_xlib,jids) {
	    set str {}
	    lappend str $jid
	    foreach col $columns {
		lappend str $stat($current_xlib,$jid,$col)
	    }
	    puts $fd [join $str "\t"]
	}
	close $fd
    }
}

proc traffic::reset {} {
    variable current_xlib
    variable columns
    variable jids
    variable stat

    if {$current_xlib != {}} return

    if {[MessageDlg .request_reset_traffic \
	     -aspect 50000 \
	     -icon warning \
	     -type user \
	     -buttons {yes no} \
	     -default 1 \
	     -cancel 1 \
	     -message [::msgcat::mc "Are you sure you want to reset counters?"]]} {
	return
    }

    # clear stats
    foreach jid $jids($current_xlib,jids) {
	foreach col $columns {
	    unset stat($current_xlib,$jid,$col)
	}
	unset jids($current_xlib,$jid)
    }
    set jids($current_xlib,jids) {}

    # clear table
    clear_and_refresh
}

# ############################
# Register in the main menu

proc traffic::setup_menu {} {
    catch {
	set m [.mainframe getmenu plugins]
	set ind [$m index end]

	$m insert $ind command -label [::msgcat::mc "Traffic statistics"] \
	    -command [list [namespace current]::open_window]
    }
}

proc traffic::desetup_menu {} {
    catch {
	set m [.mainframe getmenu plugins]
	set ind [$m index [::msgcat::mc "Traffic statistics"]]

	$m delete $ind
    }
}

##############################################################################

proc traffic::restore_window {args} {
    open_window
}

proc traffic::save_session {vsession} {
    upvar 2 $vsession session
    global usetabbar

    # We don't need JID at all, so make it empty (special case)
    set user     ""
    set server   ""
    set resource ""

    # TODO
    if {!$usetabbar} return

    set prio 0
    foreach page [.nb pages] {
	set path [ifacetk::nbpath $page]

	if {[string equal $path .traffic_stats]} {
	    lappend session [list $prio $user $server $resource \
		[list [namespace current]::restore_window] \
	    ]
	}
	incr prio
    }
}

