import csv
import logging
import os
import statistics
import time
from dataclasses import dataclass
from pathlib import Path
from subprocess import Popen

from test_automation_suite.utility.cmd import run_cmd

logger = logging.getLogger(__name__)


@dataclass
class NvidiaSmiData:
    power_draw: float
    """Actual power draw"""
    power_max: float
    """Max possible power draw"""
    temperature: float
    """Actual temperature"""


@dataclass
class NvidiaStatistics:
    power_draw_median: float
    power_draw_1_percent_low: float
    power_draw_1_percent_high: float


class Nvidia:
    @staticmethod
    def call_nvidia_smi():
        # https://gist.github.com/omerfsen/8ecb620675525ac724a92bdf5a31a4b3
        # https://docs.nvidia.com/deploy/nvidia-smi/index.html
        # nvidia-smi --help-query-gpu
        try:
            result = run_cmd(
                [
                    "nvidia-smi",
                    "--query-gpu=power.draw,power.max_limit,temperature.gpu",
                    "--format=csv,noheader,nounits",
                ]
            )
        except FileNotFoundError as error:
            raise RuntimeError("nvidia-smi is not installed or is not an Nvidia device") from error

        if not result.success:
            raise RuntimeError(f"nvidia-smi failed with status {result.status}: {result.stderr}")

        return result

    @staticmethod
    def get_nvidia_data() -> NvidiaSmiData:
        result = Nvidia.call_nvidia_smi()
        power_draw, power_max, temperature = result.stdout.strip().split(", ")

        return NvidiaSmiData(
            power_draw=float(power_draw), power_max=float(power_max), temperature=float(temperature)
        )

    @staticmethod
    def nvidia_env() -> dict[str, str]:
        """Return environment variables for NVIDIA PRIME render offloading
        Prioritizes processes to launch with the NVIDIA as preferred GPU OpenGL and Vulkan

        __NV_PRIME_RENDER_OFFLOAD=1 -> Reorders the available GPUs so that NVIDA GPU comes first for Vulkan
        __VK_LAYER_NV_optimus=NVIDIA_only -> for Vulkan make the implicit overlay for offloading to NVIDIA remove all non NVIDIA GPUs from the availability list
        __GLX_VENDOR_LIBRARY_NAME=nvidia -> OpenGL: Use NVIDIA's OpenGL implementation
        DRI_PRIME=1 -> For PRIME devices with discrete AMD GPU or if Nouveau is used, this variable should be sufficient

        https://www.tuxedocomputers.com/de/Der-umfassende-PRIME-GPU-Render-Offloading/GPU-on-demand-Mode-Leitfaden.tuxedo
        For whatever reason furmark needs DRI_PRIME
        """
        env = os.environ.copy()
        env["__NV_PRIME_RENDER_OFFLOAD"] = "1"
        env["__GLX_VENDOR_LIBRARY_NAME"] = "nvidia"
        env["__VK_LAYER_NV_optimus"] = "NVIDIA_only"
        env["DRI_PRIME"] = "1"
        return env

    # func create csv inside test-automation-suite/ time stamp folder/ nvidia
    @staticmethod
    def log_nvidia(
        csv_path: Path,
        duration: int | None = None,
        interval: float = 1.0,
        process: Popen[str] | None = None,
    ) -> None:
        file_exists = csv_path.is_file()
        with open(csv_path, "a", newline="") as nvidia_csv:
            writer = csv.writer(nvidia_csv)

            if not file_exists:
                writer.writerow(["power_draw", "power_max", "temperature"])

            end_time = time.time() + duration if duration is not None else None

            while True:
                if end_time is not None and time.time() >= end_time:
                    logger.info("Stopping NVIDIA logging: duration reached.")
                    break

                if process is not None and process.poll() is not None:
                    logger.info(
                        "Stopping NVIDIA logging: benchmark process exited with code %s.",
                        process.returncode,
                    )
                    break

                data = Nvidia.get_nvidia_data()

                writer.writerow(
                    [
                        data.power_draw,
                        data.power_max,
                        data.temperature,
                    ]
                )

                time.sleep(interval)

    @staticmethod
    def calculate_statistics(csv_file: str | Path) -> NvidiaStatistics:
        power_draw_values: list[float] = []

        with open(csv_file, newline="") as file:
            reader = csv.DictReader(file)

            for row in reader:
                value = row.get("power_draw")
                if not value:
                    continue

                power_draw_values.append(float(value))

        if not power_draw_values:
            raise ValueError(f"No power draw values found in {csv_file}")

        sorted_values = sorted(power_draw_values)

        index_1_percent = max(0, int(len(sorted_values) * 0.01) - 1)

        return NvidiaStatistics(
            power_draw_median=statistics.median(power_draw_values),
            power_draw_1_percent_low=sorted_values[index_1_percent],
            power_draw_1_percent_high=sorted_values[-index_1_percent],
        )
