import argparse
import statistics
from statistics import StatisticsError
from lemonade.state import State
from lemonade.cache import Keys
from lemonade.tools.llamacpp.utils import LlamaCppAdapter
from lemonade.tools.bench import Bench


class LlamaCppBench(Bench):
    """
    Benchmark a llama.cpp model
    """

    unique_name = "llamacpp-bench"

    def __init__(self):
        super().__init__()

        # Additional statistics generated by this bench tool
        self.status_stats.insert(
            self.status_stats.index(Keys.TOKEN_GENERATION_TOKENS_PER_SECOND) + 1,
            Keys.STD_DEV_TOKENS_PER_SECOND,
        )
        self.std_dev_token_generation_tokens_per_second_list = []

    @staticmethod
    def parser(add_help: bool = True) -> argparse.ArgumentParser:
        parser = __class__.helpful_parser(
            short_description="Benchmark an LLM in llama.cpp",
            add_help=add_help,
        )

        parser = Bench.parser(parser)

        return parser

    def run_prompt(
        self,
        state: State,
        report_progress_fn,
        prompt: str,
        iterations: int,
        warmup_iterations: int,
        output_tokens: int,
    ) -> State:
        """
        Benchmark llama.cpp model that was loaded by LoadLlamaCpp.
        """

        if self.first_run_prompt:

            if not hasattr(state, "model") or not isinstance(
                state.model, LlamaCppAdapter
            ):
                raise Exception(
                    f"{self.__class__.unique_name} requires a LlamaCppAdapter model to be "
                    "loaded first. Please run load-llama-cpp before this tool."
                )
        model: LlamaCppAdapter = state.model

        per_iteration_tokens_per_second = []
        per_iteration_time_to_first_token = []

        for iteration in range(iterations + warmup_iterations):
            try:
                # Use the adapter's generate method which already has the timeout
                # and error handling
                model.time_to_first_token = None
                model.tokens_per_second = None
                raw_output, stderr = model.generate(
                    prompt, max_new_tokens=output_tokens, return_raw=True
                )

                if model.time_to_first_token is None or model.tokens_per_second is None:
                    error_msg = (
                        "Could not find timing information in llama.cpp output.\n"
                    )
                    error_msg += "Raw output:\n" + raw_output + "\n"
                    error_msg += "Stderr:\n" + stderr
                    raise Exception(error_msg)

                self.tokens_out_len_list.append(model.response_tokens)

                if iteration > warmup_iterations - 1:
                    per_iteration_tokens_per_second.append(model.tokens_per_second)
                    per_iteration_time_to_first_token.append(model.time_to_first_token)

                report_progress_fn((iteration + 1) / (warmup_iterations + iterations))

            except Exception as e:
                error_msg = f"Failed to run benchmark: {str(e)}"
                raise Exception(error_msg)

        self.input_ids_len_list.append(model.prompt_tokens)
        mean_time_to_first_token = statistics.mean(per_iteration_time_to_first_token)
        self.mean_time_to_first_token_list.append(mean_time_to_first_token)
        self.prefill_tokens_per_second_list.append(
            model.prompt_tokens / mean_time_to_first_token
        )
        self.token_generation_tokens_per_second_list.append(
            statistics.mean(per_iteration_tokens_per_second)
        )
        try:
            self.std_dev_time_to_first_token_list.append(
                statistics.stdev(per_iteration_time_to_first_token)
            )
        except StatisticsError:
            # Less than 2 measurements
            self.std_dev_time_to_first_token_list.append(None)
        try:
            self.std_dev_token_generation_tokens_per_second_list.append(
                statistics.stdev(per_iteration_tokens_per_second)
            )
        except StatisticsError:
            # Less than 2 measurements
            self.std_dev_token_generation_tokens_per_second_list.append(None)

    def save_stats(self, state):
        super().save_stats(state)

        # Save additional statistics
        if not all(
            element is None
            for element in self.std_dev_token_generation_tokens_per_second_list
        ):
            state.save_stat(
                Keys.STD_DEV_TOKENS_PER_SECOND,
                self.get_item_or_list(
                    self.std_dev_token_generation_tokens_per_second_list
                ),
            )


# This file was originally licensed under Apache 2.0. It has been modified.
# Modifications Copyright (c) 2025 AMD
