#!/usr/bin/env python
import argparse
import models.executor
from models.docker import docker
from models.sys import sys
from models.agent import agent
from iexceptions import ExecuteException, \
    MountDirNotFoundException, InspectDockerError
from libs.unit import to_byte
from const import START, STOP, STATUS, DEFAULT_DISK_RATE


class DiskExecutor(models.executor.Executor):

    def __init__(self):
        pass

    def full(self, operation, **kwargs):

        dirname = kwargs["dirname"] if "dirname" in kwargs and kwargs[
            "dirname"] else None
        if not dirname:
            return
        container_id = kwargs["container_id"] if "container_id" in kwargs \
            and kwargs["container_id"] else None
        if not container_id:
            return
        mount_dir = docker.get_mount_dir(container_id, dirname)
        if operation == START:
            agent.set_full(mount_dir)
        else:
            agent.set_empty(mount_dir)

    def fail(self, operation, **kwargs):
        """disk use limit iops for failed"""
        dirname = kwargs["dirname"] if "dirname" in kwargs and kwargs[
            "dirname"] else None
        if not dirname:
            return
        container_id = kwargs["container_id"] if "container_id" in kwargs \
            and kwargs["container_id"] else None
        # rate = 1
        if operation == STOP:
            rate = "0"
        else:
            rate = "1"
        cgroup_path = docker.get_blkio_cgroup_path(
            container_id) + "/" + "blkio.throttle.read_bps_device"
        mount_dir = docker.get_mount_dir(container_id, dirname)
        block = sys.get_block_by_mount_in_docker(mount_dir)
        block_num = sys.get_block_number(block)
        data = block_num + " " + str(rate)
        sys.write_to_cgroup(data, cgroup_path)
        cgroup_path = docker.get_blkio_cgroup_path(
            container_id) + "/" + "blkio.throttle.write_bps_device"
        sys.write_to_cgroup(data, cgroup_path)

    def error(self, operation, **kwargs):
        """
            corrupt_files can be db/*,raft/*,db/1.sst etc
        """
        dirname = kwargs["dirname"] if "dirname" in kwargs and kwargs[
            "dirname"] else None
        if not dirname:
            return
        container_id = kwargs["container_id"] if "container_id" in kwargs \
            and kwargs["container_id"] else None
        if not container_id:
            return
        corrupt_files = kwargs["corrupt_files"] if "corrupt_files" in kwargs \
            and kwargs["corrupt_files"] else "all"
        mount_dir = docker.get_mount_dir(container_id, dirname)
        agent.write_error(mount_dir,corrupt_files)

    def limit(self, operation, **kwargs):
        dirname = kwargs["dirname"] if "dirname" in kwargs and kwargs[
            "dirname"] else None
        if not dirname:
            return
        container_id = kwargs["container_id"] if "container_id" in kwargs \
            and kwargs["container_id"] else None
        rate = kwargs["rate"] if "rate" in kwargs and kwargs[
            "rate"] else DEFAULT_DISK_RATE
        if operation == STOP:
            rate = "0"
        rate = to_byte(str(rate))
        cgroup_path = docker.get_blkio_cgroup_path(
            container_id) + "/" + "blkio.throttle.read_bps_device"
        mount_dir = docker.get_mount_dir(container_id, dirname)
        block = sys.get_block_by_mount_in_docker(mount_dir)
        block_num = sys.get_block_number(block)
        data = block_num + " " + str(rate)
        sys.write_to_cgroup(data, cgroup_path)
        cgroup_path = docker.get_blkio_cgroup_path(
            container_id) + "/" + "blkio.throttle.write_bps_device"
        sys.write_to_cgroup(data, cgroup_path)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Network Injection Simulation')
    parser.add_argument(
        'operation',
        metavar='start/stop',
        help='operations')

    parser.add_argument(
        '-a',
        '--action',
        dest='action',
        metavar='fail/full/limit/error',
        help='which action to take ')

    parser.add_argument(
        '-d',
        '--dirname',
        dest='dirname',
        help='which dirname to operation')

    parser.add_argument(
        '-r',
        '--rate',
        dest='rate',
        help='how much rate to operation, support K/M/G, default is B')

    parser.add_argument(
        '-c',
        '--container_id',
        dest='container_id',
        help='which container id to operation ')

    parser.add_argument(
        '--corrupt_files',
        dest='corrupt_files',
        help='corrupt_files to manipulation')

    args = parser.parse_args()
    executor = DiskExecutor()
    try:
        getattr(executor, args.action)(
            args.operation,
            action=args.action,
            dirname=args.dirname,
            rate=args.rate,
            container_id=args.container_id)
    except (ExecuteException,
            MountDirNotFoundException,
            InspectDockerError) as e:
        print e._msg
        exit(1)