import time
from logging import Logger
from typing import override

from test_automation_suite.benchmarks.benchmark import Benchmark
from test_automation_suite.core.context import Context
from test_automation_suite.core.task import Step, StepResult
from test_automation_suite.utility.nvidia_utility import Nvidia


class BenchmarkGpuUsageIdle(Benchmark):
    pre_benchmark_duration: int = 60
    """Waiting duration for the GPU to calm down in seconds"""

    benchmark_duration: int = 3600
    """Duration for the idle GPU measurement in seconds"""

    @override
    def steps(self) -> list[Step]:
        return [*super().steps(), self.step_run]

    def step_run(self, context: Context, logger: Logger) -> StepResult:

        # The benchmark depends on nvidia-smi for GPU monitoring, so end the check early if it is unavailable.
        try:
            Nvidia.get_nvidia_data()
        except RuntimeError as error:
            logger.error("%s", error)
            return StepResult(
                success=False,
                harvest_data={
                    "error": str(error),
                },
            )

        gpu_log_folder = context.output_directory / "gpu_log"
        gpu_log_folder.mkdir(parents=True, exist_ok=True)

        logger.info("Waiting before idle GPU measurement")
        time.sleep(self.pre_benchmark_duration)

        logger.info("Starting idle GPU measurement")
        csv_path = gpu_log_folder / f"{self.__class__.__name__}.csv"
        Nvidia.log_nvidia(
            csv_path,
            duration=self.benchmark_duration,
        )
        logger.info("Finished idle GPU measurement")
        logger.debug("NVIDIA CSV path: %s", csv_path)

        nvidia_stats = Nvidia.calculate_statistics(csv_path)
        stats = Nvidia.get_nvidia_data()

        return StepResult(
            success=True,
            harvest_data={
                "power_draw": stats.power_draw,
                "power_max": stats.power_max,
                "power_draw_median": nvidia_stats.power_draw_median,
                "power_draw_1_percent_low": nvidia_stats.power_draw_1_percent_low,
            },
        )
