#!/usr/bin/env python


import sys
import json
import time
from collections import OrderedDict

import click
from netkit.box import Box
from netkit.contrib.tcp_client import TcpClient

import peony
from peony.share import constants


class PeonyCtl(object):

    address_uri = None

    timeout = None
    username = None
    password = None

    tcp_client = None

    def __init__(self, address_uri, timeout, username, password):
        self.address_uri = address_uri
        self.timeout = timeout
        self.username = username
        self.password = password

    def make_send_box(self, cmd, username, password, payload=None):
        return Box(dict(
            cmd=cmd,
            body=json.dumps(
                dict(
                    auth=dict(
                        username=username,
                        password=password,
                    ),
                    payload=payload,
                )
            )
        ))

    def output(self, s):
        print(s)

    def start(self):

        address = self._parse_address_uri(self.address_uri)

        self.tcp_client = TcpClient(Box, address=address, timeout=self.timeout)

        try:
            self.tcp_client.connect()
        except Exception as e:
            self.output('connect fail: %s' % e)
            return False

        return True

    def handle_stat(self, loop, diff):
        """
        :param loop:
        :return:
        """
        last_result = None
        loop_times = 0

        while True:

            result = self._get_stat_once()

            if not result:
                break

            if diff:
                if last_result is not None:
                    output_result = self._diff_dicts(last_result, result)
                else:
                    output_result = None

                last_result = result
            else:
                output_result = result

            if output_result is not None:
                self.output(json.dumps(self._sort_stat_dict(output_result), indent=4))

            loop_times += 1
            if loop_times >= loop > 0:
                # 之所以不打印分割线，是为了让统计的读取工具更容易写。
                # 之前用上了正则表达式，很慢
                break

            # 如果还有下一个数据的话
            self.output('-' * 80)

            try:
                time.sleep(1)
            except KeyboardInterrupt:
                break

    def handle_stop(self):
        send_box = self.make_send_box(
            constants.CMD_ADMIN_STOP,
            self.username, self.password,
        )
        self.tcp_client.write(send_box)

        rsp_box = self.tcp_client.read()

        if not rsp_box:
            self.output('disconnected.')
            return False

        if rsp_box.ret != 0:
            self.output('fail. rsp_box.ret=%s' % rsp_box.ret)
            return False

        self.output('succ.')

    def _parse_address_uri(self, uri):
        """
        解析uri为可用的address
        :param uri: tcp://127.0.0.1:5555, /data/release/ipc.sock
        :return: address
        """

        if uri.startswith('tcp://'):
            # 文件
            uri = uri.replace('tcp://', '')

            host, port = uri.split(':')
            port = int(port)

            return (host, port)
        else:
            return uri

    def _sort_stat_dict(self, body_dict):
        """
        对统计的结果进行排序
        """
        output_items = []
        for key in ('clients', 'workers', 'pending_tasks',
                    'client_req', 'worker_req', 'discard_tasks'):

            stat_data = body_dict.get(key)

            # 因为原始数据里面没有数据的总和，需要在工具这边处理一下
            if isinstance(stat_data, dict):
                sub_items = sorted(stat_data.items(), key=lambda x: int(x[0]))
                sub_items.insert(0, ('all', sum(stat_data.values())))

                stat_data = OrderedDict(sub_items)

            output_items.append((key, stat_data))

        # OrderedDict在通过json打印的时候，会保持原来的顺序
        return OrderedDict(output_items)

    def _get_stat_once(self):
        send_box = self.make_send_box(constants.CMD_ADMIN_SERVER_STAT, self.username, self.password)
        self.tcp_client.write(send_box)

        rsp_box = self.tcp_client.read()

        if not rsp_box:
            self.output('disconnected.')
            return False

        if rsp_box.ret != 0:
            self.output('fail. rsp_box.ret=%s' % rsp_box.ret)
            return False

        return json.loads(rsp_box.body)

    def _diff_dicts(self, old_dict, new_dict):
        """
        对两个dict的数据进行差值处理
        """

        result_dict = dict()

        key_list = set(old_dict.keys()) | set(new_dict.keys())

        for key in key_list:
            if key not in old_dict:
                # 一定在new_dict里
                value = new_dict[key]
                if isinstance(value, dict):
                    result_dict[key] = self._diff_dicts(dict(), value)
                else:
                    result_dict[key] = value

            elif key not in new_dict:
                # 一定在old_dict里
                value = old_dict[key]
                if isinstance(value, dict):
                    result_dict[key] = self._diff_dicts(value, dict())
                else:
                    result_dict[key] = -value

            else:
                old_value = old_dict[key]
                new_value = new_dict[key]

                assert type(old_value) == type(new_value)

                if isinstance(old_value, dict):
                    result_dict[key] = self._diff_dicts(old_value, new_value)
                else:
                    result_dict[key] = new_value - old_value

        return result_dict


@click.group()
def cli():
    pass


@cli.command()
def version():
    """
    版本号
    """
    print(peony.__version__)


@cli.command()
@click.option('-a', '--address', default='admin.sock',
              help='peony admin address. admin.sock or tcp://127.0.0.1:9910')
@click.option('-o', '--timeout', type=int, help='connect/send/receive timeout', default=10)
@click.option('-u', '--username', help='username', default=None)
@click.option('-p', '--password', help='password', default=None)
@click.option('--loop', help='loop times, <=0 means infinite loop', type=int, default=-1)
@click.option('--diff', help='show diff values between 1 seconds', is_flag=True, default=False)
def stat(address, timeout, username, password, loop, diff):
    """
    查看统计
    """
    ctl = PeonyCtl(address, timeout, username, password)
    if not ctl.start():
        return
    ctl.handle_stat(loop, diff)


@cli.command()
@click.option('-a', '--address', default='admin.sock',
              help='peony admin address. admin.sock or tcp://127.0.0.1:9910')
@click.option('-o', '--timeout', type=int, help='connect/send/receive timeout', default=10)
@click.option('-u', '--username', help='username', default=None)
@click.option('-p', '--password', help='password', default=None)
def stop(address, timeout, username, password):
    """
    优雅停止整个服务
    """
    ctl = PeonyCtl(address, timeout, username, password)
    if not ctl.start():
        return
    ctl.handle_stop()


if __name__ == '__main__':
    cli()
