"""Low-overhead process and device sampling for performance experiments."""

from __future__ import annotations

import csv
from dataclasses import dataclass
import os
from pathlib import Path
import re
import time
from typing import Any, TextIO


_SAMPLE_FIELDS = (
	"epoch",
	"elapsed_seconds",
	"process_cpu_percent",
	"system_cpu_percent",
	"rss_mib",
	"read_mib_per_second",
	"write_mib_per_second",
	"voluntary_context_switches_per_second",
	"nonvoluntary_context_switches_per_second",
	"gpu_busy_percent",
	"vram_busy_percent",
	"gpu_sclk_mhz",
	"gpu_mclk_mhz",
	"cpu_average_mhz",
	"cpu_minimum_mhz",
	"cpu_maximum_mhz",
)


@dataclass(frozen=True)
class _Counters:
	process_ticks: int
	rss_bytes: int
	read_bytes: int
	write_bytes: int
	voluntary_context_switches: int
	nonvoluntary_context_switches: int
	system_total_ticks: int
	system_idle_ticks: int


class PerformanceSampler:
	"""Sample one game process and the active AMD GPU at a fixed interval."""

	def __init__(
		self,
		pid: int,
		output: Path,
		*,
		interval_seconds: float = 0.5,
	) -> None:
		if pid <= 0:
			raise ValueError("pid must be positive")
		if interval_seconds <= 0:
			raise ValueError("interval_seconds must be positive")
		self.pid = pid
		self.output = output
		self.interval_seconds = interval_seconds
		self._interval_ns = int(interval_seconds * 1_000_000_000)
		self._started_ns = time.time_ns()
		self._last_sample_ns: int | None = None
		self._previous: _Counters | None = None
		self._handle: TextIO = output.open("w", encoding="utf-8", newline="")
		self._writer = csv.DictWriter(self._handle, fieldnames=_SAMPLE_FIELDS)
		self._writer.writeheader()
		self._gpu_device = _discover_gpu_device()
		self._sample_count = 0
		self._peaks: dict[str, float] = {}
		self._first_counters: _Counters | None = None
		self._last_counters: _Counters | None = None

	def maybe_sample(self, now_ns: int | None = None) -> bool:
		now_ns = now_ns if now_ns is not None else time.time_ns()
		previous_sample_ns = self._last_sample_ns
		if (
			previous_sample_ns is not None
			and now_ns - previous_sample_ns < self._interval_ns
		):
			return False
		self._last_sample_ns = now_ns
		try:
			counters = _read_counters(self.pid)
		except (OSError, ValueError):
			return False
		if self._first_counters is None:
			self._first_counters = counters
		self._last_counters = counters
		row = self._sample_row(now_ns, previous_sample_ns, counters)
		self._writer.writerow(row)
		self._handle.flush()
		self._sample_count += 1
		for key, value in row.items():
			if key in {"epoch", "elapsed_seconds"} or not isinstance(value, (int, float)):
				continue
			self._peaks[key] = max(float(value), self._peaks.get(key, float("-inf")))
		self._previous = counters
		return True

	def close(self) -> dict[str, Any]:
		if not self._handle.closed:
			self._handle.close()
		first = self._first_counters
		last = self._last_counters
		return {
			"pid": self.pid,
			"output": str(self.output),
			"interval_seconds": self.interval_seconds,
			"sample_count": self._sample_count,
			"gpu_device": str(self._gpu_device) if self._gpu_device else None,
			"peaks": dict(sorted(self._peaks.items())),
			"process_io_bytes": {
				"read": max(0, last.read_bytes - first.read_bytes) if first and last else None,
				"write": max(0, last.write_bytes - first.write_bytes) if first and last else None,
			},
		}

	def _sample_row(
		self,
		now_ns: int,
		previous_sample_ns: int | None,
		counters: _Counters,
	) -> dict[str, float | str | None]:
		elapsed_seconds = (now_ns - self._started_ns) / 1_000_000_000
		row: dict[str, float | str | None] = {
			"epoch": now_ns / 1_000_000_000,
			"elapsed_seconds": elapsed_seconds,
			"process_cpu_percent": None,
			"system_cpu_percent": None,
			"rss_mib": counters.rss_bytes / 1024**2,
			"read_mib_per_second": None,
			"write_mib_per_second": None,
			"voluntary_context_switches_per_second": None,
			"nonvoluntary_context_switches_per_second": None,
			"gpu_busy_percent": _read_float(self._gpu_device, "gpu_busy_percent"),
			"vram_busy_percent": _read_float(self._gpu_device, "mem_busy_percent"),
			"gpu_sclk_mhz": _read_active_mhz(self._gpu_device, "pp_dpm_sclk"),
			"gpu_mclk_mhz": _read_active_mhz(self._gpu_device, "pp_dpm_mclk"),
			"cpu_average_mhz": None,
			"cpu_minimum_mhz": None,
			"cpu_maximum_mhz": None,
		}
		frequencies = _read_cpu_frequencies_mhz()
		if frequencies:
			row["cpu_average_mhz"] = sum(frequencies) / len(frequencies)
			row["cpu_minimum_mhz"] = min(frequencies)
			row["cpu_maximum_mhz"] = max(frequencies)
		previous = self._previous
		if previous is None or previous_sample_ns is None:
			return row
		delta_seconds = max((now_ns - previous_sample_ns) / 1_000_000_000, 0.001)
		clock_ticks = os.sysconf("SC_CLK_TCK")
		row["process_cpu_percent"] = (
			max(0, counters.process_ticks - previous.process_ticks)
			/ clock_ticks
			/ delta_seconds
			* 100
		)
		system_total = counters.system_total_ticks - previous.system_total_ticks
		system_idle = counters.system_idle_ticks - previous.system_idle_ticks
		if system_total > 0:
			row["system_cpu_percent"] = max(0, system_total - system_idle) / system_total * 100
		row["read_mib_per_second"] = (
			max(0, counters.read_bytes - previous.read_bytes) / 1024**2 / delta_seconds
		)
		row["write_mib_per_second"] = (
			max(0, counters.write_bytes - previous.write_bytes) / 1024**2 / delta_seconds
		)
		row["voluntary_context_switches_per_second"] = (
			max(
				0,
				counters.voluntary_context_switches - previous.voluntary_context_switches,
			)
			/ delta_seconds
		)
		row["nonvoluntary_context_switches_per_second"] = (
			max(
				0,
				counters.nonvoluntary_context_switches
				- previous.nonvoluntary_context_switches,
			)
			/ delta_seconds
		)
		return row


def _read_counters(pid: int) -> _Counters:
	root = Path("/proc") / str(pid)
	process_ticks, rss_bytes = _parse_proc_stat((root / "stat").read_text(encoding="utf-8"))
	io = _parse_key_values((root / "io").read_text(encoding="utf-8"))
	status = _parse_key_values((root / "status").read_text(encoding="utf-8"))
	system_total, system_idle = _parse_system_cpu(Path("/proc/stat").read_text(encoding="utf-8"))
	return _Counters(
		process_ticks=process_ticks,
		rss_bytes=rss_bytes,
		read_bytes=io.get("read_bytes", 0),
		write_bytes=io.get("write_bytes", 0),
		voluntary_context_switches=status.get("voluntary_ctxt_switches", 0),
		nonvoluntary_context_switches=status.get("nonvoluntary_ctxt_switches", 0),
		system_total_ticks=system_total,
		system_idle_ticks=system_idle,
	)


def _parse_proc_stat(value: str) -> tuple[int, int]:
	closing = value.rfind(")")
	if closing < 0:
		raise ValueError("Malformed /proc process stat")
	fields = value[closing + 2 :].split()
	if len(fields) < 22:
		raise ValueError("Truncated /proc process stat")
	process_ticks = int(fields[11]) + int(fields[12])
	rss_bytes = int(fields[21]) * os.sysconf("SC_PAGE_SIZE")
	return process_ticks, rss_bytes


def _parse_key_values(value: str) -> dict[str, int]:
	result: dict[str, int] = {}
	for line in value.splitlines():
		key, separator, raw = line.partition(":")
		if not separator:
			continue
		match = re.match(r"\s*([0-9]+)", raw)
		if match:
			result[key] = int(match.group(1))
	return result


def _parse_system_cpu(value: str) -> tuple[int, int]:
	first = value.splitlines()[0].split()
	if not first or first[0] != "cpu" or len(first) < 5:
		raise ValueError("Malformed /proc/stat CPU line")
	values = [int(item) for item in first[1:]]
	return sum(values), values[3] + (values[4] if len(values) > 4 else 0)


def _discover_gpu_device() -> Path | None:
	for path in sorted(Path("/sys/class/drm").glob("card*/device")):
		if (path / "gpu_busy_percent").is_file():
			return path
	return None


def _read_float(root: Path | None, name: str) -> float | None:
	if root is None:
		return None
	try:
		return float((root / name).read_text(encoding="utf-8").strip())
	except (OSError, ValueError):
		return None


def _read_active_mhz(root: Path | None, name: str) -> float | None:
	if root is None:
		return None
	try:
		value = (root / name).read_text(encoding="utf-8")
	except OSError:
		return None
	return _parse_active_mhz(value)


def _parse_active_mhz(value: str) -> float | None:
	match = re.search(r"([0-9]+)\s*Mhz\s*\*", value, flags=re.IGNORECASE)
	return float(match.group(1)) if match else None


def _read_cpu_frequencies_mhz() -> list[float]:
	result: list[float] = []
	for path in sorted(Path("/sys/devices/system/cpu").glob("cpu[0-9]*/cpufreq/scaling_cur_freq")):
		try:
			result.append(float(path.read_text(encoding="utf-8").strip()) / 1000)
		except (OSError, ValueError):
			continue
	return result