# Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/benchmark_serving.py

import argparse
import asyncio
import json
import logging
import os
import sys
import time
from collections import defaultdict
from datetime import datetime
from typing import Optional

import aiohttp
import pandas
import vllm
from tqdm.asyncio import tqdm
from transformers import PreTrainedTokenizerBase
from vllm.benchmarks.datasets import (
    AIMODataset,
    ASRDataset,
    BenchmarkDataset,
    BurstGPTDataset,
    ConversationDataset,
    CustomDataset,
    HuggingFaceDataset,
    InstructCoderDataset,
    MTBenchDataset,
    NextEditPredictionDataset,
    RandomDataset,
    SampleRequest,
    ShareGPTDataset,
    SonnetDataset,
    VisionArenaDataset,
)
from vllm.benchmarks.endpoint_request_func import (
    ASYNC_REQUEST_FUNCS,
    RequestFuncInput,
)
from vllm.benchmarks.serve import (
    BenchmarkMetrics,
    add_cli_args,
    calculate_metrics,
    get_tokenizer,
)

logger = logging.getLogger(__name__)
REQUEST_FUC = None


class TraceReplayDataset(BenchmarkDataset):
    # Default values copied from benchmark_serving.py for the random dataset.
    DEFAULT_PREFIX_LEN = 0
    DEFAULT_RANGE_RATIO = 0.0
    DEFAULT_INPUT_LEN = 1024
    DEFAULT_OUTPUT_LEN = 128
    REG_GROUPS = defaultdict(list)

    def load_trace(self, trace_file):
        with open(trace_file, "r", encoding="utf-8") as f:
            # Read by line
            for line in f:
                record = json.loads(line)
                self.REG_GROUPS[int(record["timestamp"]) / 1000].append(record)
        print(f"Done load trace file, time: {time.time()}")

    def __init__(
        self,
        trace_path,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        # Load trace file
        self.load_trace(trace_path)
        trace_directory_name = os.path.dirname(trace_path)
        trace_file_name = os.path.basename(trace_path).split(".")[0]
        # Filename which prompts generated by trace file would save to
        self.prompts_file_name = (
            f"{trace_directory_name}/{trace_file_name}_dataset.jsonl"
        )
        self.hash_to_tokens: dict[int, list[int]] = {}

    def generate_prompt(
        self, hash_ids: list[int], target_length: int, tokenizer
    ) -> str:
        DEFAULT_BLOCK_SIZE = 512
        vocab_size = tokenizer.vocab_size

        # Use hash_ids to influence token generation
        base_offset = hash_ids[0] if hash_ids else 0

        token_ids = []

        for i, value in enumerate(hash_ids):
            if value in self.hash_to_tokens:
                token_ids.extend(self.hash_to_tokens[value])
            elif (i + 1) * DEFAULT_BLOCK_SIZE <= target_length:
                for j in range(DEFAULT_BLOCK_SIZE):
                    token_idx = i * DEFAULT_BLOCK_SIZE + j
                    token_id = (
                        base_offset
                        + token_idx
                        + sum(
                            hash_ids[
                                token_idx % len(hash_ids) : token_idx % len(hash_ids)
                                + 3
                            ]
                        )
                    ) % vocab_size

                    self.hash_to_tokens.setdefault(value, []).append(token_id)
                token_ids.extend(self.hash_to_tokens[value])
            else:
                needed = target_length - i * DEFAULT_BLOCK_SIZE
                padding = [
                    (base_offset + len(token_ids) + j) % vocab_size
                    for j in range(needed)
                ]
                token_ids.extend(padding)

        # Decode to text and re-encode to get actual length
        text = tokenizer.decode(token_ids, skip_special_tokens=True)
        final_tokens = tokenizer.encode(text, add_special_tokens=False)

        # Truncate or pad to target length
        if len(final_tokens) > target_length:
            final_tokens = final_tokens[:target_length]
        elif len(final_tokens) < target_length:
            # Pad with more deterministic tokens
            needed = target_length - len(final_tokens)
            padding = [
                (base_offset + len(final_tokens) + i) % vocab_size
                for i in range(needed)
            ]
            final_tokens.extend(padding)

        # Final decode
        return tokenizer.decode(final_tokens, skip_special_tokens=True)

    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        save_prompts: bool = False,
        **kwargs,
    ) -> dict[float, list[SampleRequest]]:
        requests = defaultdict(list)
        # If exists prompts file generated before, load prompts from file
        if os.path.exists(self.prompts_file_name):
            with open(self.prompts_file_name, "r", encoding="utf-8") as f:
                # Read by line
                for line in f:
                    record = json.loads(line)
                    timestamp = float(record["timestamp"])
                    prompt = record["prompt"]
                    input_len = record["input_length"]
                    output_length = record["output_length"]
                    requests[timestamp].append(
                        SampleRequest(
                            prompt=prompt,
                            prompt_len=input_len,
                            expected_output_len=output_length,
                        )
                    )
            print(f"Done load prompts file, time: {time.time()}")
            return requests

        assert self.REG_GROUPS is not None, "Find no trace info!!!"

        for timestamp, record_list in self.REG_GROUPS.items():
            for req in record_list:
                hash_ids = req["hash_ids"]
                input_length = req["input_length"]
                output_length = req["output_length"]

                prompt = self.generate_prompt(hash_ids, input_length, tokenizer)
                requests[timestamp].append(
                    SampleRequest(
                        prompt=prompt,
                        prompt_len=input_length,
                        expected_output_len=output_length,
                    )
                )
        if not save_prompts:
            return requests
        with open(self.prompts_file_name, "a", encoding="utf-8") as f:
            for timestamp, reqs_list in requests.items():
                for req in reqs_list:
                    data = {
                        "timestamp": timestamp,
                        "prompt": req.prompt,
                        "input_length": req.prompt_len,
                        "output_length": req.expected_output_len,
                    }
                    f.write(json.dumps(data, ensure_ascii=False) + "\n")

        print(f"Done sample, time: {time.time()}")
        return requests


# Generate a request by benchmark dataset
def gene_one_req(
    req: json,
    tokenizer: PreTrainedTokenizerBase,
    args: argparse.Namespace,
    num_prompts=1,
):
    backend = args.backend

    if args.dataset_name == "custom":
        dataset = CustomDataset(dataset_path=args.dataset_path)
        input_requests = dataset.sample(
            num_requests=num_prompts,
            tokenizer=tokenizer,
            output_len=req["output_length"],
            skip_chat_template=args.custom_skip_chat_template,
        )

    elif args.dataset_name == "sonnet":
        dataset = SonnetDataset(dataset_path=args.dataset_path)
        # For the "sonnet" dataset, formatting depends on the backend.
        if args.backend == "openai-chat":
            input_requests = dataset.sample(
                num_requests=num_prompts,
                input_len=req["input_length"],
                output_len=req["output_length"],
                prefix_len=args.sonnet_prefix_len,
                tokenizer=tokenizer,
                return_prompt_formatted=False,
            )
        else:
            assert (
                tokenizer.chat_template or tokenizer.default_chat_template
            ), "Tokenizer/model must have chat template for sonnet dataset."
            input_requests = dataset.sample(
                num_requests=num_prompts,
                input_len=req["input_length"],
                output_len=req["output_length"],
                prefix_len=args.sonnet_prefix_len,
                tokenizer=tokenizer,
                return_prompt_formatted=True,
            )

    elif args.dataset_name == "hf":
        # all following datasets are implemented from the
        # HuggingFaceDataset base class
        if args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS:
            dataset_class = VisionArenaDataset
            args.hf_split = "train"
            args.hf_subset = None
        elif args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS:
            dataset_class = InstructCoderDataset
            args.hf_split = "train"
        elif args.dataset_path in MTBenchDataset.SUPPORTED_DATASET_PATHS:
            dataset_class = MTBenchDataset
            args.hf_split = "train"
        elif args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS:
            dataset_class = ConversationDataset
        elif args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS:
            dataset_class = AIMODataset
            args.hf_split = "train"
        elif (
            args.dataset_path in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS
        ):  # noqa: E501
            dataset_class = NextEditPredictionDataset
            args.hf_split = "train"
        elif args.dataset_path in ASRDataset.SUPPORTED_DATASET_PATHS:
            dataset_class = ASRDataset
            args.hf_split = "train"
        else:
            supported_datasets = set(
                [
                    dataset_name
                    for cls in HuggingFaceDataset.__subclasses__()
                    for dataset_name in cls.SUPPORTED_DATASET_PATHS
                ]
            )
            raise ValueError(
                f"Unsupported dataset path: {args.dataset_path}. "
                "Huggingface dataset only supports dataset_path"
                f" from one of following: {supported_datasets}. "
                "Please consider contributing if you would "
                "like to add support for additional dataset formats."
            )

        if dataset_class.IS_MULTIMODAL and backend not in [
            "openai-chat",
            "openai-audio",
        ]:
            # multi-modal benchmark is only available on OpenAI Chat backend.
            raise ValueError(
                "Multi-modal content is only supported on 'openai-chat' and "
                "'openai-audio' backend."
            )
        input_requests = dataset_class(
            dataset_path=args.dataset_path,
            dataset_subset=args.hf_subset,
            dataset_split=args.hf_split,
            random_seed=args.seed,
        ).sample(
            num_requests=num_prompts,
            tokenizer=tokenizer,
            output_len=req["output_length"],
        )

    else:
        # For datasets that follow a similar structure, use a mapping.
        dataset_mapping = {
            "sharegpt": lambda: ShareGPTDataset(
                random_seed=args.seed, dataset_path=args.dataset_path
            ).sample(
                tokenizer=tokenizer,
                num_requests=num_prompts,
                output_len=req["output_length"],
            ),
            "burstgpt": lambda: BurstGPTDataset(
                random_seed=args.seed, dataset_path=args.dataset_path
            ).sample(tokenizer=tokenizer, num_requests=num_prompts),
            "random": lambda: RandomDataset(dataset_path=args.dataset_path).sample(
                tokenizer=tokenizer,
                num_requests=num_prompts,
                prefix_len=args.random_prefix_len,
                input_len=req["input_length"],
                output_len=req["output_length"],
                range_ratio=args.random_range_ratio,
            ),
        }
        try:
            input_requests = dataset_mapping[args.dataset_name]()
        except KeyError as err:
            raise ValueError(f"Unknown dataset: {args.dataset_name}") from err
    return input_requests


# Generates prompts by dataset name
def gene_prompts_by_dataset_name(
    req_groups: dict, tokenizer: PreTrainedTokenizerBase, args: argparse.Namespace
) -> dict[float, list]:
    if args.dataset_name is None:
        raise ValueError(
            "Please specify '--dataset-name' and the corresponding "
            "'--dataset-path' if required."
        )
    # {float, list[json]}
    input_requests = defaultdict(list)
    start_time = time.perf_counter()
    for sec, reqs in sorted(req_groups.items()):
        for req in reqs:
            # Try to produce prompt by benchmark datasets
            gene_req = gene_one_req(req, tokenizer, args)
            input_requests[sec].extend(gene_req)
    gene_prompts_duration = time.perf_counter() - start_time
    print(f"Take {gene_prompts_duration} to produce prompts.")
    return input_requests


def save_metrics_to_file(
    metrics: BenchmarkMetrics, metric_percentiles: list[float], output_dir: str = "./"
):
    output_path = output_dir
    if not os.path.exists(output_path):
        os.makedirs(output_path, exist_ok=True)
    excel_file = os.path.join(output_path, "metrics.xlsx")

    outputs = {}
    outputs["time"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    outputs["total_input_tokens"] = metrics.total_input
    outputs["total_output_tokens"] = metrics.total_output
    outputs["completed requests"] = metrics.completed
    outputs["request_throughput"] = metrics.request_throughput
    outputs["request_goodput"] = metrics.request_goodput
    outputs["output_throughput"] = metrics.output_throughput
    outputs["total_token_throughput"] = metrics.total_token_throughput
    outputs["mean_ttft_ms"] = metrics.mean_ttft_ms
    for p, value in metrics.percentiles_ttft_ms:
        outputs[f"p{int(p)}_ttft_ms"] = value
    outputs["mean_tpot_ms"] = metrics.mean_tpot_ms
    for p, value in metrics.percentiles_tpot_ms:
        outputs[f"p{int(p)}_tpot_ms"] = value
    outputs["mean_itl_ms"] = metrics.mean_itl_ms
    for p, value in metrics.percentiles_itl_ms:
        outputs[f"p{int(p)}_itl_ms"] = value
    outputs["mean_e2el_ms"] = metrics.mean_e2el_ms
    for p, value in metrics.percentiles_e2el_ms:
        outputs[f"p{int(p)}_e2el_ms"] = value

    df = pandas.DataFrame([outputs])
    os.makedirs(os.path.dirname(excel_file), exist_ok=True)
    with pandas.ExcelWriter(
        excel_file, engine="openpyxl", mode="a", if_sheet_exists="replace"
    ) as writer:
        df.to_excel(writer, index=False, sheet_name="Metrics")
    print(f"Successfully saved performance metrics to {excel_file}")


def save_req_results_to_file(outputs, output_dir="./"):
    output_path = output_dir
    if not os.path.exists(output_path):
        os.makedirs(output_path, exist_ok=True)
    excel_file = os.path.join(output_path, "metrics.xlsx")
    rows = []
    # print(f"outputs: {outputs}")
    for output in outputs:
        ttft = output.ttft * 1000 if output.ttft is not None else None
        output_len = output.output_tokens if output.output_tokens is not None else 0
        input_len = output.prompt_len if output.prompt_len is not None else 0
        latency = output.latency * 1000 if output.latency is not None else None
        tpot = None
        if output_len > 1 and output.ttft is not None and output.latency is not None:
            tpot = (output.latency - output.ttft) / (output_len - 1) * 1000
        row = {
            "input_lens": input_len,
            "output_lens": output_len,
            "ttfts_ms": ttft,
            "tpot_ms": tpot,
        }
        if output.send_time and output.running_time:
            row["send_to_funning"] = output.running_time - output.send_time
        if output.running_time and output.worker_time:
            row["running_to_worker"] = output.worker_time - output.running_time
        if output.worker_time and output.start_loadkv_time:
            row["worker_to_loadkv"] = output.start_loadkv_time - output.worker_time
        if output.start_loadkv_time and output.start_forward_time:
            row["loadkv_duration"] = (
                output.start_forward_time - output.start_loadkv_time
            )
        if output.start_forward_time and output.finish_forward_time:
            row["forward_duration"] = (
                output.finish_forward_time - output.start_forward_time
            )
        if output.finish_forward_time and output.finish_savekv_time:
            row["savekv_duration"] = (
                output.finish_savekv_time - output.finish_forward_time
            )
        if output.first_token_time and output.running_time:
            row["running_to_first_token"] = (
                output.first_token_time - output.running_time
            )
        row["success"] = output.success
        rows.append(row)

    df = pandas.DataFrame(rows)
    os.makedirs(os.path.dirname(excel_file), exist_ok=True)
    with pandas.ExcelWriter(
        excel_file, engine="openpyxl", mode="a", if_sheet_exists="replace"
    ) as writer:
        df.to_excel(writer, index=False, sheet_name="details")


async def request_func(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
    session: aiohttp.ClientSession = None,
):
    if session:
        return await REQUEST_FUC(
            request_func_input=request_func_input, pbar=pbar, session=session
        )
    return await REQUEST_FUC(request_func_input=request_func_input, pbar=pbar)


# Send requests by timestamp
async def replay_trace_by_time(
    req_groups: dict, tokenizer: PreTrainedTokenizerBase, args: argparse.Namespace
):
    backend = args.backend
    model_id = args.model
    model_name = args.served_model_name
    disable_tqdm = args.disable_tqdm
    if backend not in ASYNC_REQUEST_FUNCS:
        raise ValueError(f"Unknown backend: {backend}")

    if args.base_url is not None:
        api_url = f"{args.base_url}{args.endpoint}"
        base_url = f"{args.base_url}"
    else:
        api_url = f"http://{args.host}:{args.port}{args.endpoint}"
        base_url = f"http://{args.host}:{args.port}"

    print("Starting initial single prompt test run...")
    test_request = None
    for _, reqs in sorted(req_groups.items()):
        if reqs:
            test_request = reqs[0]
            break
    if test_request is None:
        raise ValueError("No request found for initial test run.")

    use_session = tuple(map(int, vllm.__version__.split(".")[:3])) >= (0, 10, 1)
    session = None
    if use_session:
        session = vllm.aiohttp.ClientSession(
            base_url=base_url,
            trust_env=True,
            timeout=aiohttp.ClientTimeout(total=6 * 60 * 60),
        )
    global REQUEST_FUC
    REQUEST_FUC = ASYNC_REQUEST_FUNCS[backend]

    test_input = RequestFuncInput(
        model=model_id,
        model_name=model_name,
        prompt=test_request.prompt,
        api_url=api_url,
        prompt_len=test_request.prompt_len,
        output_len=test_request.expected_output_len,
        logprobs=None,
        multi_modal_content=getattr(test_request, "multi_modal_data", None),
        ignore_eos=True,
        extra_body={"temperature": 0.9},
    )

    test_output = await request_func(request_func_input=test_input, session=session)

    if not getattr(test_output, "success", False):
        raise ValueError(
            "Initial test run failed - Please make sure arguments "
            f"are correctly specified. Error: {getattr(test_output, 'error', '')}"
        )
    else:
        print("Initial test run completed. Starting main run...")

    total = sum(len(req_list) for req_list in req_groups.values())
    pbar = None if disable_tqdm else tqdm(total=total)
    semaphore = (
        asyncio.Semaphore(args.max_concurrency)
        if getattr(args, "max_concurrency", None)
        else None
    )
    start_time = time.perf_counter()
    print(f"Start time is {start_time}")
    tasks = []
    flat_requests = []

    async def _run_one_request(sample_req):
        sampling_params = {"temperature": 0.9}
        req_input = RequestFuncInput(
            model=model_id,
            model_name=model_name,
            prompt=sample_req.prompt,
            api_url=api_url,
            prompt_len=sample_req.prompt_len,
            output_len=sample_req.expected_output_len,
            logprobs=None,
            extra_body=sampling_params,
            ignore_eos=True,
        )
        if semaphore is not None:
            async with semaphore:
                return await request_func(
                    request_func_input=req_input, session=session, pbar=pbar
                )
        else:
            return await request_func(
                request_func_input=req_input, session=session, pbar=pbar
            )

    for sec, reqs in sorted(req_groups.items()):
        delay = sec - (time.perf_counter() - start_time)
        delay = max(0, delay)

        async def send_group(r=reqs, d=delay):
            await asyncio.sleep(d)
            print(
                f"Sending request at {time.perf_counter() - start_time:.3f}s with {len(r)} reqs"
            )
            group_tasks = [asyncio.create_task(_run_one_request(req)) for req in r]
            try:
                return await asyncio.gather(*group_tasks)
            except asyncio.TimeoutError:
                print(f"Request timed out: group at delay {d:.3f}s")
                return []
            except Exception as e:
                print(f"Request failed: {e}")
                return []

        tasks.append(asyncio.create_task(send_group(reqs, delay)))
        flat_requests.extend(reqs)

    group_results = await asyncio.gather(*tasks)

    if pbar is not None:
        pbar.close()
    if use_session:
        await session.close()

    outputs = []
    for res in group_results:
        if isinstance(res, list):
            outputs.extend(res)

    percentile_metrics = (
        args.metric_percentiles.split(",")
        if args.metric_percentiles
        else ["ttft", "tpot", "itl", "e2el"]
    )
    metric_percentiles = (
        [int(x) for x in args.metric_percentiles.split(",")]
        if args.metric_percentiles
        else [50, 90]
    )
    goodput = {k: float(v) for item in args.goodput for k, v in [item.split(":", 1)]}

    benchmark_duration = time.perf_counter() - start_time
    metrics, actual_output_lens = calculate_metrics(
        input_requests=flat_requests,
        outputs=outputs,
        dur_s=benchmark_duration,
        tokenizer=tokenizer,
        selected_percentiles=metric_percentiles,
        goodput_config_dict=goodput,
    )

    print("{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
    print("{:<40} {:<10}".format("Successful requests:", metrics.completed))
    print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration))
    print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))
    print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))
    print(
        "{:<40} {:<10.2f}".format(
            "Request throughput (req/s):", metrics.request_throughput
        )
    )
    print(
        "{:<40} {:<10.2f}".format(
            "Output token throughput (tok/s):", metrics.output_throughput
        )
    )
    print(
        "{:<40} {:<10.2f}".format(
            "Total Token throughput (tok/s):", metrics.total_token_throughput
        )
    )

    # Define the process_one_metric function, which can access the outer scope's selected_percentile_metrics
    def process_one_metric(
        metric_attribute_name: str,
        metric_name: str,
        metric_header: str,
    ):
        selected_percentile_metrics = percentile_metrics
        if metric_attribute_name not in selected_percentile_metrics:
            return
        print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))
        print(
            "{:<40} {:<10.2f}".format(
                f"Mean {metric_name} (ms):",
                getattr(metrics, f"mean_{metric_attribute_name}_ms"),
            )
        )
        print(
            "{:<40} {:<10.2f}".format(
                f"Median {metric_name} (ms):",
                getattr(metrics, f"median_{metric_attribute_name}_ms"),
            )
        )
        # standard deviation
        print(
            "{:<40} {:<10.2f}".format(
                f"Std {metric_name} (ms):",
                getattr(metrics, f"std_{metric_attribute_name}_ms"),
            )
        )
        for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}_ms"):
            p_word = str(int(p)) if int(p) == p else str(p)
            print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", value))

    process_one_metric("ttft", "TTFT", "Time to First Token")
    process_one_metric("tpot", "TPOT", "Time per Output Token (excl. 1st token)")
    process_one_metric("itl", "ITL", "Inter-token Latency")
    process_one_metric("e2el", "E2EL", "End-to-end Latency")
    print("=" * 50)

    output_dir = args.result_dir if args.result_dir is not None else "./"
    if args.save_result:
        save_metrics_to_file(
            metrics=metrics,
            metric_percentiles=metric_percentiles,
            output_dir=output_dir,
        )
        save_req_results_to_file(outputs=outputs, output_dir=output_dir)
    return


def create_argument_trace():
    parser = argparse.ArgumentParser(description="Benchmark LLM serving performance")
    add_cli_args(parser)

    trace_group = parser.add_argument_group("tracing parameters")
    trace_group.add_argument(
        "--trace-path",
        type=str,
        default=None,
        help="Path to trace file path.",
    )
    trace_group.add_argument(
        "--trace-mode",
        type=str,
        default="trace",
        choices=["trace", "benchmark"],
        help="Specify the trace mode: 'trace' to replay requests from cached trace files, or 'benchmark' to generate requests dynamically using the benchmark module.",
    )
    trace_group.add_argument(
        "--save-prompts",
        action="store_true",
        help="Save generated prompts with timestamp for reuse.",
    )
    return parser


def main(args: argparse.Namespace):
    print(args)
    dataset = TraceReplayDataset(args.trace_path)

    tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
    tokenizer_mode = args.tokenizer_mode
    tokenizer = get_tokenizer(
        tokenizer_id,
        tokenizer_mode=tokenizer_mode,
        trust_remote_code=args.trust_remote_code,
    )

    # Generate prompts
    if args.trace_mode == "trace":
        input_requests = dataset.sample(
            tokenizer=tokenizer, save_prompts=args.save_prompts
        )
    else:
        input_requests = gene_prompts_by_dataset_name(
            req_groups=dataset.REG_GROUPS, tokenizer=tokenizer, args=args
        )

    # Send prompts by timestamp
    asyncio.run(replay_trace_by_time(input_requests, tokenizer, args))


if __name__ == "__main__":
    # Check openpyxl for Excel export
    try:
        import openpyxl
    except ImportError:
        print("\nMissing package: openpyxl")
        print("Please install openpyxl via pip install.\n")
        sys.exit(1)

    parser = create_argument_trace()
    args = parser.parse_args()
    main(args)
