# -*- coding: utf-8 -*-
# Author      : ShiFan
# Created Date: 2023/12/29 9:12
import asyncio
# import uvloop
import sys
from logging import getLogger
from pathlib import Path
from traceback import format_exception
from typing import Optional, Union
from uuid import uuid4

import zmq.auth
import zmq.constants as z_const


base_path = Path(__file__).parent
sys.path.append(base_path.parent.as_posix())

from gaterpc.global_settings import Settings
from gaterpc.core import (
    AsyncZAPService, Context, Worker, Service, AMajordomo,
)
from gaterpc.utils import (
    HugeData, UnixEPollEventLoopPolicy, interface, msg_pack, msg_unpack,
    run_in_executor,
    to_bytes,
)
import testSettings


curve_dir = Path(__file__).parent.joinpath("curvekey/")
if curve_dir.exists():
    g_public, g_secret = zmq.auth.load_certificate(
        curve_dir.joinpath("gate.key_secret")
    )
    cw_public, _ = zmq.auth.load_certificate(
        curve_dir.joinpath("cw.key")
    )
else:
    g_public = g_secret = b""
    cw_public = b""

Settings.configure("USER_SETTINGS", testSettings)
logger = getLogger("commands")


class TAMajordomo(AMajordomo):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.clients.add(self.identity)
        self.rw_queue = asyncio.Queue()
        self.rw_task: Optional[asyncio.Task] = None
        self.rw_replies = dict()

    async def request_worker(self):
        await asyncio.sleep(3)
        try:
            rpc_service = self.services[Settings.SERVICE_DEFAULT_NAME]
        except KeyError:
            logger.error("RPC service don't running.")
            await asyncio.sleep(1)
            self.rw_task = self._loop.create_task(self.request_worker())
            return
        st = 0
        i = 0
        while 1:
            if i > 100000:
                break
            try:
                worker = await rpc_service.acquire_idle_worker()
            except asyncio.TimeoutError:
                continue
            if not st:
                st = self._loop.time()
                logger.debug(f"start time: {st}")
            log = {
                    "gtid": i + 1,
                    "action": "update",
                    "key": "tttxxxhhh",
                    "value": {
                        "name": "hostname",
                        "ip": "1.1.1.1",
                        "stat": "running" if i % 2 else "stop",
                        "remote_hosts": ["hostname1", "hostname2", "hostname3"]
                    }
                }
            request_id = to_bytes(uuid4().hex)
            body = (
                msg_pack("emit"),
                msg_pack((log,))
            )
            option = (self.identity, b"", request_id)

            # 要注意并发太多时，对套接字类型的选择和缓冲区的配置，同时适时的让出io
            # 可以适当提高Settings里的ZMQ_SOCk配置里的 z_const.HWM，
            # sysctl -w net.core.wmem_default=33554432
            # sysctl -w net.core.rmem_default=33554432
            # sysctl -w net.core.wmem_max=67108864
            # sysctl -w net.core.rmem_max=67108864

            # 可以顺序等待发送
            # await self.send_to_backend(
            #     worker.identity, Settings.MDP_COMMAND_REQUEST,
            #     option, body
            # )
            # 也可以创建为发送任务，乱序发送
            self.backend_tasks.add(
                t := self._loop.create_task(
                    self.send_to_backend(
                            worker.identity, Settings.MDP_COMMAND_REQUEST,
                            option, body
                        )
                )
            )
            t.add_done_callback(self.cleanup_backend_task)
            if log["gtid"] == 100000:
                logger.debug(f"sent use time: {self._loop.time() - st}")
            self.rw_replies[request_id] = self._loop.create_future()
            i += 1
            await asyncio.sleep(0)

    rw_replies_num = 0

    async def reply_frontend(
        self, client_id: bytes,
        service_name: Union[str, bytes],
        request_id: Union[str, bytes],
        *body
    ):
        if client_id == self.identity:
            if request_id in self.rw_replies:
                # logger.debug(f"result: {msg_unpack(body[0])}")
                self.rw_replies.pop(request_id).set_result(body)
                if not self.rw_replies_num:
                    logger.debug(f"first replies: {self._loop.time()}")
                self.rw_replies_num += 1
                if not self.rw_replies_num % 10000:
                    logger.debug(
                        f"worker replies number: {self.rw_replies_num}, "
                        f"time: {self._loop.time()}"
                    )
        else:
            await super().reply_frontend(
                client_id, service_name, request_id, *body
            )

    async def test_rw(self):
        await asyncio.sleep(5)
        logger.info("start test rw")
        try:
            st = self._loop.time()
            for i in range(100000):
                log = {
                    "gtid": i + 1,
                    "action": "update",
                    "key": "tttxxxhhh",
                    "value": {
                        "name": "hostname",
                        "ip": "1.1.1.1",
                        "stat": "running" if i % 2 else "stop",
                        "remote_hosts": ["hostname1", "hostname2", "hostname3"]
                    }
                }
                await self.rw_queue.put(log)
                await asyncio.sleep(0)
            et = self._loop.time()
            logger.info(f"put use time: {et - st}")
        except Exception as e:
            logger.info("*****************************************************")
            for line in format_exception(*sys.exc_info()):
                logger.error(line)
            logger.info("*****************************************************")
            raise e

    def run(self):
        if not self._broker_task:
            self._broker_task = self._loop.create_task(self._broker_loop())
        if not self.rw_task:
            self.rw_task = self._loop.create_task(self.request_worker())

    def stop(self):
        if self.rw_task:
            if not self.rw_task.cancelled():
                self.rw_task.cancel()
            self.rw_task = None
        if self._broker_task:
            if not self._broker_task.cancelled():
                self._broker_task.cancel()
            self.close()
            self._broker_task = None


async def zap_server(ctx):
    zap = AsyncZAPService(context=ctx)
    zap.configure_plain(
        Settings.ZAP_DEFAULT_DOMAIN,
        {
            Settings.ZAP_PLAIN_DEFAULT_USER: Settings.ZAP_PLAIN_DEFAULT_PASSWORD
        }
    )
    try:
        zap.start()
        await zap._recv_task
    finally:
        zap.stop()


def start_zap(ctx):
    loop = asyncio.new_event_loop()
    loop.slow_callback_duration = 0.01
    loop.run_until_complete(zap_server(ctx))


async def majordomo(
    frontend=None, backend=None, bind_gate=None, connect_gate=None
):
    Settings.setup()
    # loop = asyncio.get_event_loop()
    # loop.slow_callback_duration = 0.01
    if backend:
        Settings.WORKER_ADDR = backend
    Settings.ZAP_REPLY_TIMEOUT = 10.0
    ctx = Context()
    # zipc = Settings.ZAP_ADDR
    # logger.info(f"zap ipc addr: {zipc}")
    # zap = AsyncZAPService(addr=zipc)
    # zap.configure_plain(
    #     Settings.ZAP_DEFAULT_DOMAIN,
    #     {
    #         Settings.ZAP_PLAIN_DEFAULT_USER: Settings.ZAP_PLAIN_DEFAULT_PASSWORD
    #     }
    # )

    gr_majordomo = TAMajordomo(
        context=ctx,
        # gate_zap_mechanism=Settings.ZAP_MECHANISM_PLAIN,
        # gate_zap_credentials=(
        #     Settings.ZAP_PLAIN_DEFAULT_USER,
        #     Settings.ZAP_PLAIN_DEFAULT_PASSWORD
        # )
    )
    if g_secret:
        gr_majordomo.bind_backend(sock_opt={
                z_const.CURVE_SECRETKEY: g_secret,
                z_const.CURVE_PUBLICKEY: g_public,
                z_const.CURVE_SERVER: True,
            })
    else:
        gr_majordomo.bind_backend()
    if frontend:
        if g_secret:
            gr_majordomo.bind_frontend(
                frontend,
                sock_opt={
                    z_const.CURVE_SECRETKEY: g_secret,
                    z_const.CURVE_PUBLICKEY: g_public,
                    z_const.CURVE_SERVER: True,
                }
            )
        else:
            gr_majordomo.bind_frontend(frontend)
    if bind_gate:
        gr_majordomo.bind_gate(bind_gate)
    await asyncio.sleep(1)
    logger.info("start test")
    try:
        # zap.start()
        # gr_majordomo.connect_zap(zap_addr=zipc)
        gr_majordomo.run()
        if connect_gate:
            try:
                await gr_majordomo.connect_gate(connect_gate)
            except Exception as e:
                logger.error(e)
                raise e
        await asyncio.sleep(5)
        # await gr_majordomo.test_rw()
        await gr_majordomo.rw_task
        await gr_majordomo._broker_task
    finally:
        exc = gr_majordomo.rw_task.exception()
        gr_majordomo.stop()
        logger.info(f"request_zap.cache_info: {gr_majordomo.zap_cache}")
        # zap.stop()
        if exc:
            raise exc


def test(
    frontend="ipc:///tmp/gate-rpc/run/c1",
    backend=None, bind_gate=None, connect_gate=None
):
    print(f"frontend: {frontend}, backend: {backend}")
    print(f"bind_gate: {bind_gate}, connect_gate: {connect_gate}")
    asyncio.set_event_loop_policy(UnixEPollEventLoopPolicy())
    asyncio.run(
        majordomo(frontend, backend, bind_gate, connect_gate)
    )


if __name__ == "__main__":
    argv = sys.argv[1:]
    test(*argv)
