# -*- coding: utf-8 -*-
"""
Display transfer rates of a tor instance.

Configuration parameters:
    cache_timeout: An integer specifying the cache life-time of the modules
        output in seconds (default 2)
    control_address: The address on which the Tor daemon listens for control
        connections (default "127.0.0.1")
    control_password: The password to use for the Tor control connection
        (default None)
    control_port: The port on which the Tor daemon listens for control
        connections (default 9051)
    format: A string describing the output format for the module
        (default "↑ {up} ↓ {down}")
    format_value: A string describing how to format the transfer rates
        (default "[\?min_length=12 {rate:.1f} {unit}]")
    rate_unit: The unit to use for the transfer rates
        (default "B/s")
    si_units: A boolean value selecting whether or not to use SI units
        (default False)

Format placeholders:
    {down} The incoming transfer rate
    {up} The outgoing transfer rate

format_value placeholders:
    {rate} The current transfer-rate's value
    {unit} The current transfer-rate's unit

Requires:
    stem: python module from pypi https://pypi.python.org/pypi/stem

Example:

```
tor_rate {
    cache_timeout = 10
    format = "IN: {down} | OUT: {up}"
    control_port = 1337
    control_password = "TertiaryAdjunctOfUnimatrix01"
    si_units = True
}

order += "tor_rate"
```

@author Felix Morgner <felix.morgner@gmail.com>
@license 3-clause-BSD

SAMPLE OUTPUT
{'full_text': u'\u2191 652.3 B/s \u2193 938.1 B/s'}
"""

from stem import ProtocolError, SocketError
from stem.connection import AuthenticationFailure
from stem.control import Controller, EventType

ERROR_AUTHENTICATION = "Error: Failed to authenticate with Tor daemon!"
ERROR_CONNECTION = "Error: Failed to establish control connection!"
ERROR_PROTOCOL = "Error: Failed to register event handler!"


class Py3status:
    """
    """

    # available configuration parameters
    cache_timeout = 2
    control_address = "127.0.0.1"
    control_password = None
    control_port = 9051
    format = u"↑ {up} ↓ {down}"
    format_value = "[\?min_length=12 {rate:.1f} {unit}]"
    rate_unit = "B/s"
    si_units = False

    def post_config_hook(self):
        self._auth_failure = False
        self._down = 0
        self._handler_active = False
        self._up = 0

    def tor_rate(self, outputs, config):
        text = ""
        if not self._handler_active and not self._auth_failure:
            try:
                self._register_event_handler()
            except ProtocolError:
                text = ERROR_PROTOCOL
            except SocketError:
                text = ERROR_CONNECTION
            except AuthenticationFailure:
                text = ERROR_AUTHENTICATION
                self._auth_failure = True
        elif self._auth_failure:
            text = ERROR_AUTHENTICATION
        else:
            text = self.py3.safe_format(self.format, self._get_rates())

        return {"cached_until": self.py3.time_in(self.cache_timeout), "full_text": text}

    def _get_rates(self):
        up, up_unit = self.py3.format_units(
            self._up, unit=self.rate_unit, si=self.si_units
        )
        down, down_unit = self.py3.format_units(
            self._down, unit=self.rate_unit, si=self.si_units
        )
        return {
            "up": self.py3.safe_format(
                self.format_value, {"rate": up, "unit": up_unit}
            ),
            "down": self.py3.safe_format(
                self.format_value, {"rate": down, "unit": down_unit}
            ),
        }

    def _handle_event(self, event):
        self._down = event.read
        self._up = event.written

    def _register_event_handler(self):
        self._control = Controller.from_port(
            address=self.control_address, port=self.control_port
        )
        if self.control_password:
            self._control.authenticate(password=self.control_password)
        self._control.add_event_listener(lambda e: self._handle_event(e), EventType.BW)
        self._handler_active = True


if __name__ == "__main__":
    from py3status.module_test import module_test

    config = {"control_password": "SevenOfNine"}
    module_test(Py3status, config)
