# Standard
from dataclasses import dataclass
import argparse
import asyncio
import logging
import random
import time

# Third Party
from transformers import AutoTokenizer
from utils import (
    AsyncLoopWrapper,
    PromptBuildMethodType,
    build_rag_prompt,
    compute_f1,
    compute_rl,
    init_logger,
    load_dataset,
)
import openai
import pandas as pd

logger = init_logger(__name__, logging.INFO)

system_prompt_set = {
    PromptBuildMethodType.QA: "You will be asked a question after reading several passages. "  # noqa: E501
    "Please directly answer the question based on the given passages. "
    "Do NOT repeat the question. "
    "The answer should be within 5 words..\nPassages:\n",
    PromptBuildMethodType.FEW_SHOT: "Summarize the dialogue into a few short sentences. "  # noqa: E501
    "The following are some examples.\n\n",
}
query_prompt_set = {
    PromptBuildMethodType.QA: "\n\nAnswer the question directly based on the given passages."  # noqa: E501
    " Do NOT repeat the question. "
    "The answer should be within 5 words. \nQuestion:",
    PromptBuildMethodType.FEW_SHOT: "",
}


@dataclass
class WorkloadConfig:
    # Overall QPS
    qps: float
    # Model name
    model: str
    # Tokenizer name
    tokenizer: str
    # Dataset.
    dataset: str
    # Start index of the workload
    start_index: int
    # End index of the workload
    end_index: int
    # Random shuffle.
    shuffle: bool
    # System prompt.
    system_prompt: str
    # Separator.
    separator: str
    # Query prompt.
    query_prompt: str
    # Prompt build method.
    prompt_build_method: PromptBuildMethodType
    # Max tokens for each generation.
    max_tokens: int


@dataclass
class Response:
    request_id: int
    body: str
    ttft: float
    generation_time: float
    prompt_tokens: int
    generation_tokens: int
    launch_time: float
    finish_time: float


def parse_arguments():
    parser = argparse.ArgumentParser(description="Parse RAG benchmark configurations.")
    parser.add_argument("--qps", type=float, required=True, help="Overall QPS")
    parser.add_argument("--model", type=str, required=True, help="Model name")
    parser.add_argument("--tokenizer", type=str, default="", help="Tokenizer name")
    parser.add_argument("--dataset", type=str, required=True, help="The dataset path")
    parser.add_argument(
        "--start-index", type=int, default=0, help="Start index of the workload"
    )
    parser.add_argument(
        "--end-index", type=int, default=-1, help="End index of the workload"
    )
    parser.add_argument("--shuffle", action="store_true", help="Random shuffle")
    parser.add_argument("--system-prompt", type=str, default="", help="System prompt")
    parser.add_argument("--separator", type=str, default="", help="Separator")
    parser.add_argument("--query-prompt", type=str, default="", help="Query prompt")
    parser.add_argument(
        "--prompt-build-method",
        type=str,
        required=True,
        help="Prompt build method",
    )
    parser.add_argument(
        "--base-url",
        type=str,
        required=True,
        help="Base URL of the serving engine endpoint",
    )
    parser.add_argument(
        "--api-key",
        type=str,
        default="EMPTY",
        help="API key of the serving engine endpoint",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="summary.csv",
        help="The output file name (ended with csv or txt) for the summary csv and txt",
    )
    parser.add_argument(
        "--warmup", action="store_true", help="Whether to enable warmup"
    )
    parser.add_argument(
        "--time",
        type=int,
        default=None,
        help="The total running time in seconds",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Whether to enable verbose logging",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=32,
        help="Max tokens for each generation",
    )
    parser.add_argument(
        "--step-interval", type=float, default=0.02, help="Step interval"
    )
    args = parser.parse_args()
    return args


class RequestExecutor:
    def __init__(
        self,
        base_url: str,
        api_key: str,
        prompt_build_method: PromptBuildMethodType,
        model: str,
    ):
        self.client = openai.AsyncOpenAI(api_key=api_key, base_url=base_url)
        self.model = model
        self.loop = AsyncLoopWrapper.GetOrStartLoop()
        self.prompt_build_method = prompt_build_method

    async def _async_launch_request(self, request_id, prompt, max_tokens):
        start_time = time.time()
        first_token_time = None
        words = ""
        response = None
        if self.prompt_build_method == PromptBuildMethodType.QA:
            messages = [{"role": "user", "content": prompt}]
            response = await self.client.chat.completions.create(
                messages=messages,
                model=self.model,
                temperature=0,
                stream=True,
                max_tokens=max_tokens,
                stream_options={"include_usage": True},
            )
        elif self.prompt_build_method == PromptBuildMethodType.FEW_SHOT:
            response = await self.client.completions.create(
                prompt=prompt,
                model=self.model,
                temperature=0,
                stream=True,
                max_tokens=max_tokens,
                stream_options={"include_usage": True},
            )
        else:
            raise ValueError(f"Invalid prompt build method {self.prompt_build_method}")
        async for tok in response:
            if not tok.choices:
                continue
            chunk_message = tok.choices[0].delta.content
            if chunk_message is not None:
                if first_token_time is None and chunk_message != "":
                    first_token_time = time.time()
                words += chunk_message
        tokens_out = tok.usage.completion_tokens
        tokens_prefill = tok.usage.prompt_tokens
        finish_time = time.time()
        return Response(
            request_id=request_id,
            body=words,
            ttft=first_token_time - start_time,
            generation_time=finish_time - first_token_time,
            prompt_tokens=tokens_prefill,
            generation_tokens=tokens_out,
            launch_time=start_time,
            finish_time=finish_time,
        )

    def launch_request(self, request_id: int, prompt, max_tokens, finish_callback):
        """
        finish_callback: Callable[[Response], None]
        """
        real_callback = lambda x: finish_callback(x.result())
        future = asyncio.run_coroutine_threadsafe(
            self._async_launch_request(request_id, prompt, max_tokens),
            self.loop,
        )
        future.add_done_callback(real_callback)


def warmup_engine(executor: RequestExecutor):
    logger.info("Warming up the engine")
    for i in range(10):
        prompt = f"WARMUP: Hi, I'm user {i}. Here are some text: {'hi ' * 100}."
        executor.launch_request(-1, prompt, 100, lambda x: None)

    AsyncLoopWrapper.WaitLoop()
    logger.info("Warm up finished.")


class RAGManager:
    def __init__(self, workload_config: WorkloadConfig):
        self.workload_config = workload_config
        eval_dataset = load_dataset(workload_config.dataset)
        start_index = workload_config.start_index
        end_index = workload_config.end_index
        if end_index < 0:
            end_index = len(eval_dataset)
        eval_dataset = eval_dataset[start_index:end_index]
        if workload_config.shuffle:
            random.shuffle(eval_dataset)
        self._prompts = []
        self._answers = []
        self._build_method = workload_config.prompt_build_method
        self._generated_text = []
        self._generation_time = []
        self._prefill_tok_cnt = []
        self._generation_tok_cnt = []
        self._ttft = []
        self._tpot = []
        for ex in eval_dataset:
            prompt, _ = build_rag_prompt(
                workload_config.system_prompt,
                ex,
                workload_config.query_prompt,
                workload_config.separator,
                workload_config.prompt_build_method,
            )
            self._prompts.append(prompt)
            self._answers.append(ex["answers"])
            self._generated_text.append(None)
            self._generation_time.append(None)
            self._prefill_tok_cnt.append(None)
            self._generation_tok_cnt.append(None)
            self._ttft.append(None)
            self._tpot.append(None)
        self._tokenizer = AutoTokenizer.from_pretrained(workload_config.tokenizer)
        self._last_request_time = -1.0
        self._last_request_index = 0
        assert workload_config.qps > 0
        self._gap = 1.0 / workload_config.qps
        self._max_tokens = workload_config.max_tokens

    def _update_result(self, response: Response):
        self._generated_text[response.request_id] = response.body
        self._ttft[response.request_id] = response.ttft
        self._tpot[response.request_id] = (
            response.generation_time / response.generation_tokens
        )
        self._generation_time[response.request_id] = response.generation_time
        self._prefill_tok_cnt[response.request_id] = response.prompt_tokens
        self._generation_tok_cnt[response.request_id] = response.generation_tokens

    def step(self, timestamp: float, executor: RequestExecutor) -> bool:
        if self._last_request_index >= len(self._prompts):
            return False
        if (
            self._last_request_time < 0
            or timestamp >= self._last_request_time + self._gap
        ):
            prompt = self._prompts[self._last_request_index]
            request_id = self._last_request_index
            self._last_request_time = timestamp
            self._last_request_index += 1
            executor.launch_request(
                request_id, prompt, self._max_tokens, self._update_result
            )
        return True

    def summary(self, start_time: float, end_time: float) -> pd.DataFrame:
        cnt = len(self._ttft)
        assert cnt > 0
        avg_ttft = sum(self._ttft) / cnt
        avg_tpot = sum(self._tpot) / cnt
        # Create a dataframe
        quality = []
        for i in range(cnt):
            if self._build_method == PromptBuildMethodType.QA:
                quality.append(
                    max(
                        [
                            compute_f1(self._generated_text[i], answer, self._tokenizer)
                            for answer in self._answers[i]
                        ]
                    )
                )
            elif self._build_method == PromptBuildMethodType.FEW_SHOT:
                quality.append(
                    max(
                        [
                            compute_rl(self._generated_text[i], answer)
                            for answer in self._answers[i]
                        ]
                    )
                )
            else:
                raise ValueError(f"Invalid prompt build method {self._build_method}")
        avg_quality = sum(quality) / cnt
        df = pd.DataFrame(
            {
                "quality": quality,
                "ttft": self._ttft,
                "tpot": self._tpot,
                "generation_time": self._generation_time,
                "prefill_token_cnt": self._prefill_tok_cnt,
                "generation_token_cnt": self._generation_tok_cnt,
            }
        )
        total_time = end_time - start_time
        thput = cnt / total_time
        logger.info(
            f"Summary: {cnt} requests, average_ttft={avg_ttft} (second)\n"
            f" average_tpot={avg_tpot} (second)\n"
            f"throughput={thput} (req/s)\n"
            f"average_quality={avg_quality}\n"
        )
        return df


def run_rag(args):
    build_prompt_method_str = args.prompt_build_method.upper()
    build_prompt_method = None
    if build_prompt_method_str == "QA":
        build_prompt_method = PromptBuildMethodType.QA
    elif build_prompt_method_str == "FEW_SHOT":
        build_prompt_method = PromptBuildMethodType.FEW_SHOT
    else:
        raise ValueError(f"Invalid prompt build method {build_prompt_method_str}")
    workload_config = WorkloadConfig(
        qps=args.qps,
        model=args.model,
        tokenizer=args.tokenizer,
        dataset=args.dataset,
        start_index=args.start_index,
        end_index=args.end_index,
        shuffle=args.shuffle,
        system_prompt=args.system_prompt,
        separator=args.separator,
        query_prompt=args.query_prompt,
        prompt_build_method=build_prompt_method,
        max_tokens=args.max_tokens,
    )
    executor = RequestExecutor(
        base_url=args.base_url,
        api_key=args.api_key,
        prompt_build_method=build_prompt_method,
        model=args.model,
    )
    if args.warmup:
        warmup_engine(executor)
    manager = RAGManager(workload_config)
    step_interval = args.step_interval
    num_steps = 0
    start_time = time.time()
    try:
        while True:
            num_steps += 1
            effective = manager.step(time.time(), executor)
            if not effective:
                break
            time.sleep(step_interval)
            if args.time is not None and time.time() - start_time > args.time:
                break

    except KeyboardInterrupt:
        logger.info("Interrupted, waiting for the final result")

    AsyncLoopWrapper.StopLoop()

    logger.info(f"Finished benchmarking, dumping summary to {args.output}")
    summary = manager.summary(start_time, time.time())
    summary.to_csv(args.output, index=False)


def main():
    args = parse_arguments()
    build_prompt_method_str = args.prompt_build_method.upper()
    build_prompt_method = None
    if build_prompt_method_str == "QA":
        build_prompt_method = PromptBuildMethodType.QA
    elif build_prompt_method_str == "FEW_SHOT":
        build_prompt_method = PromptBuildMethodType.FEW_SHOT
    else:
        raise ValueError(f"Invalid prompt build method {build_prompt_method_str}")
    if len(args.system_prompt) == 0:
        args.system_prompt = system_prompt_set[build_prompt_method]
    if len(args.query_prompt) == 0:
        args.query_prompt = query_prompt_set[build_prompt_method]
    if len(args.tokenizer) == 0:
        args.tokenizer = args.model
    args.system_prompt = args.system_prompt.encode().decode("unicode_escape")
    args.query_prompt = args.query_prompt.encode().decode("unicode_escape")
    if args.verbose:
        global logger
        logger = init_logger(__name__, level=logging.DEBUG)
    run_rag(args)


if __name__ == "__main__":
    main()
