import subprocess
import sys


def install_packages():
    packages = [
        "jcclang",
        "trio",
        "noiseprotocol",
        "argon2-cffi",
        "dacite",
    ]
    for pkg in packages:
        print(f"正在安装 {pkg} ...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", pkg,
            "-i", "https://pypi.tuna.tsinghua.edu.cn/simple"
        ])


# 本地调试，如果依赖已经安装，可以注释掉
install_packages()

from client import run_as_client
from leader import run_as_leader
from model import SimpleCNN
from utils import *

import os

import trio

from jcclang.core.jobhub import JobHubClient

JOBHUB_HOST = os.getenv("FL_JOBHUB_HOST", "101.201.215.196")
JOBHUB_PORT = int(os.getenv("FL_JOBHUB_PORT", "7901"))
SHARED_SECRET = os.getenv("FL_JOBHUB_SECRET", "01234567890123456789012345678901").encode()

JOBSET_ID = os.getenv("FL_JOBSET_ID", "fl_mnist")
SELF_JOB_ID = os.getenv("FL_JOB_ID", "2")
LEADER_JOB_ID = LEADER_PREFIX + SELF_JOB_ID

HEARTBEAT_INTERVAL = 5
HEARTBEAT_TIMEOUT = 15

# 数据集路径，只需要填写路径即可，不用具体到数据
LOCAL_DATA = ["./data1", "./data2"]

cnn_model = SimpleCNN()


async def leader_thread():
    async with JobHubClient.connect(JOBHUB_HOST, JOBHUB_PORT, JOBSET_ID, LEADER_JOB_ID, SHARED_SECRET) as cli:
        # 本地调试，需要注册数据集，线上需注释掉
        await tmp_init_data(cli, SELF_JOB_ID, LOCAL_DATA)

        while True:
            if not await is_leader_alive(cli) and await try_become_leader(cli=cli, job_id=LEADER_JOB_ID):
                logger.info("leader not alive, try become leader")
                await run_as_leader(job_id=LEADER_JOB_ID, jobhub_cli=cli, scheduler=None, model_template=cnn_model,
                                    initial_config_path="./initial_ownership.json", checkpoint_path="./checkpoint")

            await trio.sleep(HEARTBEAT_INTERVAL)


async def client_thread():
    async with JobHubClient.connect(JOBHUB_HOST, JOBHUB_PORT, JOBSET_ID, SELF_JOB_ID, SHARED_SECRET) as cli:
        while True:
            if await is_leader_alive(cli) and not await is_client_alive(cli, SELF_JOB_ID):
                logger.info("client is not alive, try become client")
                await run_as_client(cli, SELF_JOB_ID)

            await trio.sleep(HEARTBEAT_INTERVAL)


async def agent_main():
    # 打印所有环境变量信息
    logger.info(f"Environment Variables: {os.environ}")
    logger.info(f"FL Agent started: job_id={SELF_JOB_ID}")

    # async with JobHubClient.connect(JOBHUB_HOST, JOBHUB_PORT, JOBSET_ID, SELF_JOB_ID, SHARED_SECRET) as cli:
    async with trio.open_nursery() as nursery:
        nursery.start_soon(leader_thread)
        nursery.start_soon(client_thread)


if __name__ == "__main__":
    trio.run(agent_main)
