#!/usr/bin/env python
import argparse
import models.executor
from iexceptions import ExecuteException
from const import DEFAULT_EGRESS_BANDWIDTH, DEFAULT_INGRESS_BANDWIDTH, \
    DEFAULT_LATENCY, DEFAULT_LATENCY_VARIATION, START, STOP


class BandwidthExecutor(models.executor.Executor):
    """as bandwidth"""

    def __init__(self):
        self._wondershaper = "wondershaper"

    def common_runner(self, operation, egress, ingress, latency,
                      latency_variation, **kwargs):
        interface = kwargs[
            "interface"] if "interface" in kwargs and kwargs["interface"] else "em1"

        if egress is None:
            egress = DEFAULT_EGRESS_BANDWIDTH
        if ingress is None:
            ingress = DEFAULT_INGRESS_BANDWIDTH
        if latency is None:
            latency = DEFAULT_LATENCY
        if latency_variation is None:
            latency_variation = DEFAULT_LATENCY_VARIATION

        if operation == START:
            cmd = "{wondershaper} -a {interface} -d {ingress} " \
                "-u {egress} -l {latency} -n {latency_var}".format(
                wondershaper=self._wondershaper,
                interface=interface,
                ingress=ingress,
                egress=egress,
                latency=latency,
                latency_var=latency_variation)
        elif operation == STOP:
            cmd = "{wondershaper} -c -a {interface}".format(
                wondershaper=self._wondershaper, interface=interface)

        try:
            self._execute_or_revert_cmd(cmd)
        except ExecuteException as e:
            raise e
        return True

    def limit(self, operation, egress, ingress, latency, latency_variation,
              **kwargs):
        return self.common_runner(
            operation=args.operation,
            egress=args.egress,
            ingress=args.ingress,
            latency=args.latency,
            latency_variation=args.latency_variation,
            interface=interface)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Bandwidth Injection Simulation")
    parser.add_argument(
        '-a',
        '--action',
        dest='action',
        metavar='limit',
        help='which action to take ')
    parser.add_argument('operation', metavar="start/stop", help="operations")
    parser.add_argument('--host_ip', dest='host_ip', help="which host ip")
    parser.add_argument(
        '--egress', dest='egress', help="set maximum upload rate")
    parser.add_argument(
        '--ingress', dest='ingress', help="set maximum download rate")
    parser.add_argument('--latency', dest='latency', help="set latency")
    parser.add_argument(
        '--latency_variation',
        dest='latency_variation',
        help="latency add random variation")
    parser.add_argument(
        '-i',
        '--interface',
        dest='interface',
        help='which interface to operation ')

    args = parser.parse_args()
    executor = BandwidthExecutor()

    interface = args.interface
    if args.host_ip:
        from libs.network import get_netdev_for_host_ip
        interface = get_netdev_for_host_ip(args.host_ip)

    try:
        getattr(executor, args.action)(
            operation=args.operation,
            egress=args.egress,
            ingress=args.ingress,
            latency=args.latency,
            latency_variation=args.latency_variation,
            interface=interface)
    except ExecuteException as e:
        print e._msg
        exit(1)