#!/usr/bin/env python
#
# Written by Boudewijn Schoon. peer-to-peer@frayja.com
#

from Tribler.Connection import Progress
encode = Progress.protocol_encode
decode = Progress.protocol_decode

import SocketServer
import threading 
import time
import os
import stat
import curses
import re
import zlib
import string

def bytes_to_string(i):
    """Returns a string description of an integer in bytes.

    i.e.
    1024 -> 1k
    512  -> 512b
    """
    if i < 0:
        return (i, "-")
    if i < 1024:
        return (i, "b")
    elif i < 1048576:
        return (i / 1024, "k")
    elif i < 1073741824:
        return (i / 1048576, "m")
    else:
        return (i / 1073741824, "g")

colors = (curses.COLOR_BLACK,
          curses.COLOR_RED,
          curses.COLOR_GREEN,
          curses.COLOR_YELLOW,
          curses.COLOR_BLUE,
          curses.COLOR_MAGENTA,
          curses.COLOR_CYAN,
          curses.COLOR_WHITE)
color_iter = (0, 1, len(colors))

class DummyPeer:
    def __init__(self, peer_id=None):
        assert peer_id is None or type(peer_id) is str
        if peer_id is None:
            self._peer_id = None
        else:
            self._peer_id = peer_id

    def __getitem__(self, key):
        return None

    def __str__(self):
        return "Dummy<%s>" % re.sub("[^a-zA-Z0-9-]", "?", str(self._peer_id))

    def __iter__(self):
        return iter(())

    def get(self, key, default=None):
        return default

    def get_timestamp(self, key):
        return 0

    def get_color(self):
        return 0

class Peer(SocketServer.BaseRequestHandler):

    template_gnuplot = string.Template(
        """#!/usr/bin/gnuplot
reset
set terminal png
set key noautotitle
set key left

set title "Upload and download bandwidth for ${config_progress_client_name}"
set xlabel "Time (seconds)"
set ylabel "Transferred (megabytes)"
set output "${working_directory}/${filename}-bandwidth.png"
set yrange [-10:]
plot "${working_directory}/${filename}" using 1:($$2)/1024/1024 title "download" with lines, \\
     "${working_directory}/${filename}" using 1:($$3)/1024/1024 title "upload" with lines
system "publish p2p/${working_directory}/${filename}-bandwidth.png ${working_directory}/${filename}-bandwidth.png"
#system "rm ${filename}-bandwidth.png"

set title "Workload on system running ${config_progress_client_name}"
set xlabel "Time (seconds)"
set ylabel "Workload (using unix workload)"
set output "${working_directory}/${filename}-workload.png"
set yrange [0:2]
plot "${working_directory}/${filename}" using 1:($$6) title "workload" with lines
system "publish p2p/${working_directory}/${filename}-workload.png ${working_directory}/${filename}-workload.png"
#system "rm ${filename}-workload.png"

set title "Teaming statistics for ${config_progress_client_name}"
set xlabel "Time (seconds)"
set ylabel "Number of groups"
set output "${working_directory}/${filename}-teaming.png"
set yrange [0:10]
plot "${working_directory}/${filename}" using 1:($$7) title "supervising" with lines, \\
     "${working_directory}/${filename}" using 1:($$8) title "teaming" with lines
system "publish p2p/${working_directory}/${filename}-teaming.png ${working_directory}/${filename}-teaming.png"
#system "rm ${filename}-teaming.png"
""")
        
    _lock = threading.Lock()
    _peer_dict = {}
    _peer_list = []

    @classmethod
    def get_peer(cls, peer_id, default=None):
        assert peer_id is None or type(peer_id) is str, peer_id
        cls._lock.acquire()
        try:
            if peer_id in cls._peer_dict:
                return cls._peer_dict[peer_id]
            elif default is None:
                return DummyPeer(peer_id)
            else:
                return default
        finally:
            cls._lock.release()

    @classmethod
    def get_peers(cls):
        cls._lock.acquire()
        l = cls._peer_list[:]
        cls._lock.release()
        return l
            
    def setup(self):
        self._parsers = {"config-user-peer-id":self._parse_user_peer_id,
                         "config-progress-client-name":self._parse_progress_client_name,
                         "connection-list":self._parse_connection_list,
                         "trace":self._parse_trace}

        self._connected = True
        self._info = {}
        self._trace_list = []
        self._files = None

        global color_iter
        low, high, max = color_iter
        color_id = low * max + high + 10
        curses.init_pair(color_id, colors[low], colors[high])
        high += 1
        if high == max:
            high = 0
            low += 1

            if low == max:
                low = 0

        color_iter = (low, high, max)
        self._color = curses.color_pair(color_id)

        self._lock.acquire()
        self._peer_list.append(self)
        self._lock.release()

    def handle(self):
        buff = ""
        recv = self.request.recv
        info = self._info
        decompress = zlib.decompressobj()
        ratio = (0, 0)
        while True:

            tmp = recv(4048)
            if len(tmp) == 0:
                return
            buff += decompress.decompress(tmp)

            try:
                buff, required_length, messages = decode(buff, len(buff))
            except:
                self._info["connection-error"] = (time.time(), "Unable to decode incoming data")
                break

            now = time.time()
            for timestamp, key, value in messages:
                parser = self._parsers.get(key)
                if parser:
                    value = apply(parser, (value,))
                    
                self._info[key] = (now, value)

    def finish(self):
        self._connected = False
        if self._files:
            for f in self._files.itervalues():
                f.close()
            self._files = None

    def _parse_user_peer_id(self, value):
        self._lock.acquire()
        self._peer_dict[value] = self
        self._lock.release()
        return value

    def _parse_progress_client_name(self, value):

        if self._files:
            return value

        self._files = _files = {}

        # figure out file names that do not clash with currently
        # existing files
        i = 0
        cont = True
        while cont:
            cont = False

            filenames = {"data":"%s/%s-data-%d"%(working_directory, value, i),
                         "gnuplot":"%s/%s-gnuplot-%d"%(working_directory, value, i),
                         }

            for filename in filenames:
                if os.access(filename, os.F_OK):
                    i += 1
                    cont = True
                    break

        # open all files
        for key, filename in filenames.iteritems():
            _files[key] = open(filename, "w")

        # write the gnuplot script
        f = _files["gnuplot"]
        f.write(self.template_gnuplot.substitute({"filename":os.path.basename(filenames["data"]),
                                                  "working_directory":working_directory,
                                                  "config_progress_client_name":value}))
        f.close()
        os.chmod(f.name, stat.S_IRWXU | stat.S_IRGRP | stat.S_IROTH)
        del _files["gnuplot"]

        gnuplots_file.write("gnuplot %s\n" % f.name)

        return value

    def _parse_trace(self, value):
        try:
            last_id, last_message = self._trace_list[-1]
        except IndexError:
            last_id, last_message = (0, None)
        
        while len(self._trace_list) > 75:
            self._trace_list.pop(0)

        self._trace_list.append((last_id+1, value))

    def _parse_connection_list(self, value):
        d = {}
        for connection_id, (peer_id, address, nice) in value.iteritems():
            d[connection_id] = (Peer.get_peer(peer_id), address, nice)
        return d

    def __getitem__(self, key):
        return self._info.get(key, (None, None))[1]

    def __str__(self):
        return str(self._info.get("config-progress-client-name", (None, id(self)))[1])

    def __iter__(self):
        for key, (timestamp, value) in self._info.iteritems():
            yield timestamp, key, value

    def get(self, key, default=None):
        return self._info.get(key, (0, default))[1]

    def get_timestamp(self, key):
        return self._info.get(key, (0, None))[0]

    def get_color(self):
        return self._color

    def is_connected(self):
        return self._connected

    def get_trace_list(self):
        return self._trace_list

class ThreadingTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
    pass

class GuiCurses(threading.Thread):

    def __init__(self):
        threading.Thread.__init__(self)
        self.start()
        
    def run(self):

        stdscr = curses.initscr()
        stdscr.keypad(1)
        curses.start_color()
        curses.noecho()
        curses.cbreak()
        curses.halfdelay(20)
        curses.curs_set(0)

        try:
            start_time = time.time()
            addstr = stdscr.addstr

            displays = ((0, "trace", self._print_trace),
                        (1, "connections", self._print_connections),
                        (2, "bittorrent", self._print_bittorrent),
                        (3, "teaming", self._print_teaming),
                        (4, "text", self._print_text))

            selected_peer = 0
            selected_display = 4

            while True:
                maxy, maxx = stdscr.getmaxyx()

                # fill with spaces
                for y in range(0, maxy-1):
                    addstr(y, 0, " "*(maxx))

                x = maxx
                for index, title, callback in displays:
                    if selected_display == index:
                        color = curses.A_BOLD
                    else:
                        color = 0
                    addstr(0, x-len(title), title, color)
                    x -= len(title) + 3

                peer_list = Peer.get_peers()
                try:
                    peer = peer_list[selected_peer]
                except IndexError:
                    if peer_list:
                        peer = peer_list[0]
                        selected_peer = 0
                    else:
                        peer = None

                if peer:
                    apply(displays[selected_display][2], (peer, stdscr, (len(peer_list)+4, 0, maxy, maxx)))
                    addstr(len(peer_list)+3, 0, "-"*maxx, peer.get_color())
                self._print_peers(displays[selected_display][1], peer_list, stdscr, (0, 0, len(peer_list)+3, maxx))

                self._store_progress(peer_list)

                key = stdscr.getch()

                if key in (ord("q"), ord("Q")):
                    break
                elif key == curses.KEY_UP:
                    if selected_peer == 0:
                        selected_peer = len(peer_list) - 1
                    else:
                        selected_peer -= 1
                elif key == curses.KEY_DOWN:
                    selected_peer += 1
                elif key == curses.KEY_RIGHT:
                    if selected_display == 0:
                        selected_display = len(displays) - 1
                    else:
                        selected_display -= 1
                elif key == curses.KEY_LEFT:
                    selected_display += 1
                    if selected_display == len(displays):
                        selected_display = 0                    

        finally:
            stdscr.keypad(0)
            curses.nocbreak()
            curses.echo()
            curses.endwin()

    def cleanup(self):
        pass

    def _store_progress(self, peer_list):
        for peer, f in [(peer, peer._files["data"]) for peer in peer_list if peer._files]:
            down, up, down_speed, up_speed = peer.get("connection-statistic", (-1, -1, -1, -1))
            f.write("%d %d %d %d %d %f %d %d\n" % (int(time.time()) - time_offset,
                                             down,
                                             up,
                                             down_speed,
                                             up_speed,
                                             peer.get("load-average", (-1.00,))[0],
                                             len(peer.get("teaming42-supervising-list", ())),
                                             len(peer.get("teaming42-teaming-list", ()))))

    def _print_peers(self, display, peer_list, scr, (miny, minx, maxy, maxx)):
        addstr = scr.addstr

        addstr(miny, minx+26, "down")
        addstr(miny, minx+38, "up")
        addstr(miny, minx+46, "recv")
        addstr(miny, minx+56, "send")
        addstr(miny, minx+66, "progress")
        addstr(miny, minx+77, "load")
        addstr(miny, minx+87, "connection-error")

        if display == "teaming":
            addstr(miny, minx+110, "super")
            addstr(miny, minx+117, "team")
            total_supervising = 0
            total_teaming = 0

        total_down = 0
        total_up = 0
        total_down_speed = 0
        total_up_speed = 0
        y = miny+1
        left_list = []
        for peer in peer_list:
            if peer.is_connected():
                color_mod = curses.A_BOLD
            else:
                color_mod = 0

            down, up, down_speed, up_speed = peer.get("connection-statistic", (-1, -1, -1, -1))

            # calculate the progress
            progress_list = []
            for torrent_id in peer.get("torrent-list", ()):
                pieces_downloaded, pieces_total, pieces_str = peer.get("%d-torrent-completed"%torrent_id, (1, -100, None))
                progress_list.append("%d%%"%(pieces_downloaded*100/pieces_total))
            progress = ",".join(progress_list)

            # get the average load
            load = peer.get("load-average", (-1.00,))[0]
            
            addstr(y, minx+0, "%23s"%peer, peer.get_color())
            addstr(y, minx+25, "%5d"%(down_speed/1024), color_mod)
            addstr(y, minx+35, "%5d"%(up_speed/1024), color_mod)
            addstr(y, minx+44, "%5d%s"%bytes_to_string(down), color_mod)
            addstr(y, minx+54, "%5d%s"%bytes_to_string(up), color_mod)
            addstr(y, minx+65, "%8s"%progress, color_mod)
            addstr(y, minx+77, "%s"%load, color_mod)
            addstr(y, minx+87, "%s"%peer["connection-error"], color_mod)

            if display == "teaming":
                count = len(peer.get("teaming42-supervising-list", ()))
                addstr(y, minx+110, "%5d"%count)
                total_supervising += count

                count = len(peer.get("teaming42-teaming-list", ()))
                addstr(y, minx+116, "%5d"%count)
                total_teaming += count

            y += 1
            if y >= maxy: return

            if down > 0 and peer.is_connected():
                total_down += down
                total_up += up
                total_down_speed += down_speed
                total_up_speed += up_speed

        addstr(y, minx+25, "-----")
        addstr(y, minx+35, "-----")
        addstr(y, minx+45, "-----")
        addstr(y, minx+55, "-----")
        if display == "teaming":
            addstr(y, minx+112, "---")
            addstr(y, minx+118, "---")
        y += 1
        if y >= maxy: return

        addstr(y, minx+25, "%5d"%(total_down_speed/1024))
        addstr(y, minx+35, "%5d"%(total_up_speed/1024))
        addstr(y, minx+44, "%5d%s"%bytes_to_string(total_down))
        addstr(y, minx+54, "%5d%s"%bytes_to_string(total_up))
        if display == "teaming":
            addstr(y, minx+110, "%5d"%total_supervising)
            addstr(y, minx+116, "%5d"%total_teaming)
        y += 1
        if y >= maxy: return

    def _print_connections(self, peer, scr, (miny, minx, maxy, maxx)):
        addstr = scr.addstr
        y = miny
        connection_dict = peer["connection-list"]
        if connection_dict:
            for connection_id, (connection, address, nice) in connection_dict.iteritems():
                down, up, down_speed, up_speed = peer.get("%d-connection-statistic"%connection_id, (-1, -1, -1, -1))

                addstr(y, minx, "%23s"%connection, connection.get_color())
                addstr(y, minx+25, "%5d"%(down_speed/1024))
                addstr(y, minx+35, "%5d"%(up_speed/1024))
                addstr(y, minx+44, "%5d%s"%bytes_to_string(down))
                addstr(y, minx+54, "%5d%s"%bytes_to_string(up))
                addstr(y, minx+77, "%s:%d"%address)
                addstr(y, minx+110, "%f"%nice)
                y += 1
                if y >= maxy: return
        else:
            addstr(miny, minx, "There are no connections")

    def _print_bittorrent(self, peer, scr, (miny, minx, maxy, maxx)):
        addstr = scr.addstr
        y = miny
        connection_dict = peer["connection-list"]
        if connection_dict:

            addstr(y, minx+77, "am-int")
            addstr(y, minx+87, "am-cho")
            addstr(y, minx+97, "is-int")
            addstr(y, minx+107, "is-cho")
            addstr(y, minx+117, "her-req")
            addstr(y, minx+127, "my-req")

            y += 1
            if y >= maxy: return

            for connection_id, (connection, address, nice) in connection_dict.iteritems():
                down, up, down_speed, up_speed = peer.get("%d-connection-statistic"%connection_id, (-1, -1, -1, -1))
                have_count, total_count, bit_list = peer.get("%d-bittorrent-bitfield"%connection_id, (-1, 100, ()))

                addstr(y, minx, "%23s"%connection, connection.get_color())
                addstr(y, minx+25, "%5d"%(down_speed/1024))
                addstr(y, minx+35, "%5d"%(up_speed/1024))
                addstr(y, minx+44, "%5d%s"%bytes_to_string(down))
                addstr(y, minx+54, "%5d%s"%bytes_to_string(up))
                addstr(y, minx+69, "%d"%(have_count * 100 / total_count))

                addstr(y, minx+77, "%s"%peer.get("%d-bittorrent-am-interested"%connection_id, (-1,)))
                addstr(y, minx+87, "%s"%peer.get("%d-bittorrent-am-choking"%connection_id, (-1,)))
                addstr(y, minx+97, "%s"%peer.get("%d-bittorrent-is-interested"%connection_id, (-1,)))
                addstr(y, minx+107, "%s"%peer.get("%d-bittorrent-is-choking"%connection_id, (-1,)))
                addstr(y, minx+117, "%d"%peer.get("%d-bittorrent-her-requests"%connection_id, (-1,)))
                addstr(y, minx+127, "%d"%peer.get("%d-bittorrent-my-requests"%connection_id, (-1,)))

                y += 1
                if y >= maxy: return
        else:
            addstr(miny, minx, "There are no connections")

    def _print_teaming(self, peer, scr, (miny, minx, maxy, maxx)):
        y = miny
        addstr = scr.addstr
        torrent_dict = peer["torrent-list"]
        connection_dict = peer["connection-list"]
        supervising_dict = peer["teaming42-supervising-list"]
        teaming_dict = peer["teaming42-teaming-list"]
        dummy = DummyPeer()

        if supervising_dict:
            for (torrent_id, piece_id), connection_ids in supervising_dict.iteritems():
                torrent_name = torrent_dict.get(torrent_id, (torrent_id,))[0]
                addstr(y, minx, "Supervising group %d [%s]" % (piece_id, torrent_name))
                for connection_id in connection_ids:
                    y += 1
                    if y >= maxy: return
                    connection = connection_dict.get(connection_id, (dummy,))[0]
                    addstr(y, minx+3, "%20s"%connection, connection.get_color())
                y += 2
                if y >= maxy: return

        addstr(y, minx, "-"*(maxx-minx))
        y += 2
        if y >= maxy: return

        if teaming_dict:
            for (torrent_id, piece_id), (supervisor_id, connection_ids) in teaming_dict.iteritems():
                downloaded_blocks, total_blocks = peer.get("%d-%d-teaming42-progress"%(torrent_id, piece_id), (1, -100))
                progress = "%d%%" % (downloaded_blocks*100/total_blocks)

                torrent_name = torrent_dict.get(torrent_id, (torrent_id,))[0]
                addstr(y, minx, "[%s] Team-member of team %d [%s]" % (progress, piece_id, torrent_name))
                y += 1
                if y >= maxy: return

                supervisor = connection_dict.get(supervisor_id, (dummy,))[0]
                addstr(y, minx, "Supervisor:")
                addstr(y, minx+12, "%s"%supervisor, supervisor.get_color())

                for connection_id in connection_ids:
                    y += 1
                    if y >= maxy: return
                    connection = connection_dict.get(connection_id, (dummy,))[0]
                    addstr(y, minx+3, "%20s"%connection, connection.get_color())
                y += 2
                if y >= maxy: return

        addstr(y, minx, "-"*(maxx-minx))
        y += 2
        if y >= maxy: return

        if torrent_dict:
            for torrent_id, (torrent_name, torrent_size, piece_size) in torrent_dict.iteritems():
                addstr(y, minx, "Interested [%s]"%torrent_name)
                y += 1
                if y >= maxy: return

                for connection_id, (connection, address, nice) in connection_dict.iteritems():
                    count, total, interested_list = peer.get("%d-teaming42-interested"%connection_id, (0, 0, ()))
                    l = [" "] * total
                    for bit in interested_list:
                        l[bit] = "-"
                    
                    l.insert(0, "[")
                    l.append("]")

                    addstr(y, minx+3, "%20s"%connection, connection.get_color())
                    addstr(y, minx+25, "".join(l)[:maxx-minx-25])

                    y += 1
                    if y >= maxy: return
                    

    def _print_trace(self, peer, scr, (miny, minx, maxy, maxx)):
        addstr = scr.addstr
        y = miny
        max_len = maxx - minx - 5
        for ident, message in peer.get_trace_list()[-(maxy-miny-1):]:
            addstr(y, minx, "%3d: %s" % (ident, message[:max_len]))
            y += 1

    def _print_text(self, peer, scr, (miny, minx, maxy, maxx)):
        addstr = scr.addstr
        y = miny
        now = time.time()
        ignore_count = 0
        max_len = maxx - minx - 37

        items = list(peer.__iter__())
        items.sort(reverse=True)

        for timestamp, key, value in items[:maxy-miny-1]:
            addstr(y, minx, key[:35])
            addstr(y, minx+37, str(value)[:max_len])

            if now - timestamp > 60:
                addstr(y, maxx-3, ">1m")
            else:
                addstr(y, maxx-3, "%ds" % (now - timestamp))
            y += 1

class GuiConsole(threading.Thread):

    def __init__(self):
        threading.Thread.__init__(self)
        self._continue = True
        self.start()

    def run(self):
        pass
#         start = time.time()

#         f = file("progress-%d"%time.time(), "w")
#         while self._continue:

#             l = [str(time.time() - start)]
#             down, up, down_speed, up_speed = peer.get("connection-statistic", (0, 0, 0, 0))
#             l.extend(["%d %d"%(down, up) for peer in peer_list])
#             l.extend("\n")
#             f.write(" ".join(l))

#             print " ".join(l),
#             time.sleep(5)

    def cleanup(self):
        self._continue = False

working_directory = time.strftime("./progress/%Y%m%d.%H%M%S")
os.makedirs(working_directory)

gnuplots_file = open("%s/gnuplots.sh"%working_directory, "w")
gnuplots_file.write("#!/bin/bash\n")
os.chmod(gnuplots_file.name, stat.S_IRWXU | stat.S_IRGRP | stat.S_IROTH)

time_offset = int(time.time())
        
try:
    gui = GuiCurses()
    server = ThreadingTCPServer(("", 47771), Peer)
    server.serve_forever()
except Exception, e:
    gui.cleanup()
    print e

# add shell code to generate a viewable HTML page

template = string.Template("""
HTML=${working_directory}/images.html

echo "<html><body><table>" >> $$HTML

for FILE in ${working_directory}/*.png
do
  echo "<tr><td><img alt='$${FILE}' src='http://www.frayja.com/pub/p2p/$${FILE}'></td></tr>" >> $$HTML
done

echo "</table></body></html>" >> $$HTML

publish p2p/${working_directory}/images.html $${HTML}
echo "A HTML summary can be found at:"
echo "www.frayja.com/pub/p2p/${working_directory}/images.html::open"
""")
gnuplots_file.write(template.substitute({"working_directory":working_directory}))
gnuplots_file.close()

print gnuplots_file.name