"""
RPyC Service wrapper around the engine supporting tensor parallelism.
If there are 4 GPUs, initialize 4 services. Rank 0 is the main, the rest are helpers.
Main runs the full engine and broadcasts tensors to helpers.
Start each EngineService in a separate process. Connect and init each.
Use RPyC to communicate with main worker to add jobs and poll generated tokens.
"""

import rpyc
from rpyc.utils.server import ThreadedServer
import os
import torch
import torch.multiprocessing as mp
import torch.distributed as dist
import time
import pickle
from tqdm import tqdm
import datetime

from suno_utils.gpt.engine import Engine, Request


def pickle_result(func):
    """
    Decorator to pickle the result of a function.
    """

    def wrapped(*args, **kwargs):
        return pickle.dumps(func(*args, **kwargs))

    return wrapped


def terminate_processes():
    import os
    import signal
    import psutil

    def kill_child_processes(parent_pid, sig=signal.SIGTERM):
        try:
            parent = psutil.Process(parent_pid)
        except psutil.NoSuchProcess:
            return
        children = parent.children(recursive=True)
        for process in children:
            process.send_signal(sig)

    # Get the current process ID
    pid = os.getpid()

    # Kill all child processes
    kill_child_processes(pid)


class EngineService(rpyc.Service):
    def on_connect(self, conn):
        from torch._inductor import config

        config.triton.cudagraph_trees = False

    def exposed_init(
        self,
        ckpt_path: str,
        tokenizer_path: str,
        max_sequences: int,
        compile: bool,
        min_batch_size: int,
        batch_increment: int,
        rank: int,
        world_size: int,
        max_length_s: float = None,
    ):
        self.rank = rank
        self.world_size = world_size

        os.environ["LOCAL_WORLD_SIZE"] = str(world_size)
        os.environ["LOCAL_RANK"] = str(rank)
        os.environ["ENABLE_INTRA_NODE_COMM"] = "1"

        print(f"Rank {rank} world size {world_size}, initing process group")
        torch.cuda.set_device(rank)
        # leaving backend blank inits both gloo and nccl backends. gloo is necessary for cpu data
        dist.init_process_group(
            init_method="tcp://localhost:28765",
            rank=rank,
            world_size=world_size,
            timeout=datetime.timedelta(hours=10000),  # dont timeout while waiting for requests
        )
        print(
            f"Initialized process group for rank {dist.get_rank()} and world size {dist.get_world_size()}"
        )

        dist.barrier()  # sanity check process group initialization

        self.exposed_engine = Engine(
            ckpt_path,
            tokenizer_path,
            max_sequences=max_sequences,
            compile=compile,
            min_batch_size=min_batch_size,
            batch_increment=batch_increment,
            use_tp=world_size > 1,
            max_length_s=max_length_s,
        )

        print(f"Rank {rank} world size {world_size}, starting engine")

        self.exposed_engine.start()
        self.exposed_engine.wait_for_warmup()

        print(f"Rank {rank} world size {world_size}, engine started")

    def get_engine(self):
        return self.exposed_engine

    @pickle_result
    def get_generated_codes(self, job_id: str, start=0):
        """
        Retrieve generated codes for a given job starting from a specified index.

        Args:
            job_id (str): The unique identifier for the job.
            start (int, optional): The starting index from which to retrieve the generated codes. Defaults to 0.

        Returns:
            list or None: A list of generated codes if available, an empty list if the job is not completed and no codes are available,
                          or None if the job is completed and no codes are available.
        """
        job = self.exposed_engine.jobs[job_id]
        codes = job.generated_tokens[start:]
        if len(codes) > 0:
            return codes

        if job.completed:
            return None  # signal completed with None
        else:
            return []

    @pickle_result
    def get_job(self, job_id: str):
        return self.exposed_engine.jobs[job_id]

    @pickle_result
    def get_model_cfg(self):
        return self.exposed_engine.model.config

    @pickle_result
    def get_tokenizer(self):
        return self.exposed_engine.tokenizer

    def get_free_caches(self):
        return len(self.exposed_engine.free_caches())

    def get_used_caches(self):
        return len(self.exposed_engine.used_caches())

    def add_request(self, pickled_request: bytes):
        request: Request = pickle.loads(pickled_request)
        return self.exposed_engine.add_request_blocking(request)

    def run_request(self, picked_requests: bytes, **kwargs):
        requests = pickle.loads(picked_requests)
        return self.exposed_engine.run_request(requests, **kwargs)


def init_env(port: int):
    t = ThreadedServer(
        EngineService,
        port=port,
        protocol_config={"allow_all_attrs": True, "allow_pickle": True},
    )
    t.start()


class EngineRpcClient:
    """
    A client for interacting with the EngineService via RPC.

    This client provides methods to add and run requests on the EngineService.
    It handles the serialization and deserialization of requests using pickle.
    """

    def __init__(self, engine_service: EngineService):
        self.engine_service = engine_service

    def add_request(self, request: Request):
        return self.engine_service.add_request(pickle.dumps(request))

    def run_request(self, request: Request, **kwargs):
        return self.engine_service.run_request(pickle.dumps(request), **kwargs)

    def get_generated_codes(self, job_id: str, start=0):
        return pickle.loads(self.engine_service.get_generated_codes(job_id, start))

    def get_job(self, job_id: str):
        return pickle.loads(self.engine_service.get_job(job_id))

    def get_model_cfg(self):
        return pickle.loads(self.engine_service.get_model_cfg())

    def get_tokenizer(self):
        return pickle.loads(self.engine_service.get_tokenizer())


def start_service_processes(
    ckpt_path: str,
    tokenizer_path: str,
    max_sequences: int,
    min_batch_size: int = 4,
    batch_increment: int = 8,
    max_length_s: float = None,
    compile=True,
    port: int = 18861,
    world_size: int = 1,
):
    """
    Initialize and start multiple service processes for the GPT engine.

    This function sets up and starts a specified number of service processes,
    each running on a different port. It also waits for all services to start
    and initializes/warms up the engines.

    Returns:
        EngineRpcClient: A client to interact with the initialized engine services.
    """
    processes = []
    for i in range(world_size):
        p = mp.get_context("forkserver").Process(target=init_env, args=(port + i,))
        p.start()
        processes.append(p)

    # wait for all the services to start
    for i in tqdm(range(world_size), desc="Waiting for services to start"):
        service_started = False
        while not service_started:
            try:
                c = rpyc.connect(
                    "localhost",
                    port + i,
                    config={"allow_all_attrs": True, "allow_pickle": True},
                )
                service_started = True
            except BaseException:
                print(f"Service {i} not started yet, waiting...")
                time.sleep(1)

    # connect to all the services and initialize the engines
    connections = []
    inits = []
    for i in tqdm(range(world_size), desc="Initializing engines"):
        c = rpyc.connect(
            "localhost",
            port + i,
            config={"allow_all_attrs": True, "allow_pickle": True},
        )
        async_init = rpyc.async_(c.root.init)
        init_state = async_init(
            ckpt_path,
            tokenizer_path,
            max_sequences,
            compile=compile,
            min_batch_size=min_batch_size,
            batch_increment=batch_increment,
            rank=i,
            world_size=world_size,
            max_length_s=max_length_s,
        )
        init_state.set_expiry(2 * 60)
        connections.append(c)
        inits.append(init_state)

    for i in range(world_size):
        inits[i].wait()
        if inits[i].error or inits[i].expired:
            terminate_processes()
            raise Exception(f"Rank {i} init failed: {inits[i].error}, {inits[i].value}")
        print(f"Rank {i} init done.")
    connections[0]._config["sync_request_timeout"] = None  # No timeout

    time.sleep(5)

    return EngineRpcClient(connections[0].root)