# traffic.tcl --
#
#       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 length [::xmpp::xml::toText $xmldata]]

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

            set jid [::xmpp::xml::getAttr $attrs $attr]
            if {$jid eq ""} {
                set jid SERVER
            } else {
                catch { set jid [::xmpp::jid::removeResource $jid] }
            }
        }
        default {
            set tag ""
            set size [string length $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 1m]
    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 0.25m]
    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 {[lindex $wished_xlib 0] ne $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 {[string range $slave 0 $max] ne "$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 {$type eq "off" && [lsearch $connections $wished_xlib] < 0} {
        set wished_xlib [list {} None]
    }

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

    if {$type eq "on" && [string equal $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

    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 ne ""} {
        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
    }
}

# vim:ft=tcl:ts=8:sw=4:sts=4:et
