import time
from logging import Logger
from subprocess import PIPE, Popen
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.dependencies.dependency import Dependency
from test_automation_suite.dependencies.init_apt_dependencies import (
    mangohud_dependency,
)
from test_automation_suite.dependencies.init_unpack_dependencies import (
    cyberpunk_dependency,
)
from test_automation_suite.utility.mangohud_utility import MangoHud
from test_automation_suite.utility.nvidia_utility import Nvidia


class BenchmarkEdgeGaming(Benchmark):
    @override
    def dependencies(self) -> list[Dependency]:
        return [*super().dependencies(), mangohud_dependency, cyberpunk_dependency]

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

    def step_run(self, context: Context, logger: Logger) -> StepResult:
        benchmark_name = "Cyberpunk2077"

        # 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("Prepare MangoHUD")
        hud = MangoHud(gpu_log_folder / benchmark_name)

        logger.info("Running Cyberpunk benchmark")
        process = BenchmarkEdgeGaming.start_CP2077()

        logger.info("Logging NVIDIA data while Cyberpunk benchmark is running")
        # Delay the Nvidia data collection to make sure it starts after cyberpunk has open
        time.sleep(10)
        csv_path = gpu_log_folder / f"{self.__class__.__name__}.csv"
        Nvidia.log_nvidia(
            csv_path,
            process=process,
        )
        return_code = process.wait()
        if return_code != 0:
            logger.error("Cyberpunk exited with error code %s", return_code)

            return StepResult(success=False)

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

        fps_stats = hud.calculate_fps_statistics()

        success = fps_stats.fps_median > 0.0 and nvidia_stats.power_draw_median > 0.0
        if not success:
            logger.error(
                f"FPS statistics ({fps_stats.fps_median}) or power draw ({nvidia_stats.power_draw_median}) were 0"
            )

        return StepResult(
            success=success,
            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,
                "fps_avg": fps_stats.fps_avg,
                "fps_median": fps_stats.fps_median,
                "fps_min": fps_stats.fps_min,
                "fps_max": fps_stats.fps_max,
                "fps_1_percent_low": fps_stats.fps_1_percent_low,
            },
        )

    @staticmethod
    def start_CP2077() -> Popen[str]:
        cp2077_path = cyberpunk_dependency.installation_path

        cmd = f"""
        cd "{cp2077_path}" && \
        STEAM_COMPAT_CLIENT_INSTALL_PATH="{cp2077_path}/Steam" \
        STEAM_COMPAT_DATA_PATH="{cp2077_path}/Prefixes/default/Cyberpunk 2077" \
        __NV_PRIME_RENDER_OFFLOAD="1" \
        __GLX_VENDOR_LIBRARY_NAME="nvidia" \
        __VK_LAYER_NV_optimus="NVIDIA_only" \
        mangohud "{cp2077_path}/GE-Proton-latest/proton" run \
        "{cp2077_path}/Cyberpunk 2077/bin/x64/Cyberpunk2077.exe" \
        --launcher-skip -benchmark -skipStartScreen
        """

        return Popen(
            cmd,
            shell=True,
            stdout=PIPE,
            stderr=PIPE,
            text=True,
        )
