"""音声やファイル名を含まない処理計測。測定失敗は本処理へ伝播させない。"""
import functools
import inspect
import json
import math
import re
import subprocess
import threading
import time
from pathlib import Path

_lock = threading.Lock()


def _path(root, jid):
    if not re.fullmatch(r"[a-f0-9]{32}", str(jid)):
        raise ValueError("invalid_job_id")
    return Path(root) / jid / "gpu_usage.json"


def _update(root, jid, update):
    try:
        with _lock:
            target = _path(root, jid)
            try:
                data = json.loads(target.read_text(encoding="utf-8"))
            except (OSError, ValueError):
                data = {}
            update(data)
            temporary = target.with_suffix(".usage.tmp")
            temporary.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
            temporary.replace(target)
    except Exception:
        print("[GPU Usage] measurement_write_failed", flush=True)


def begin(root, jid, duration=None):
    def change(data):
        data.clear()
        data.update(schema_version=1, request_started_at=int(time.time()*1000),
                    processing_started_at=int(time.time()*1000),
                    audio_duration_seconds=duration if isinstance(duration, (int, float)) and math.isfinite(duration) and duration > 0 else None,
                    processing_ms=None, queue_wait_ms=None, gpu_seconds=None,
                    gpu_seconds_is_estimated=False, gpu_function_ms=None,
                    device_utilization_pct=None, gpu_sample_count=0, status="processing")
    _update(root, jid, change)


def read(root, jid):
    try:
        with _lock:
            return json.loads(_path(root, jid).read_text(encoding="utf-8"))
    except Exception:
        return None


def finish(root, jid, succeeded):
    def change(data):
        end = int(time.time()*1000)
        data.update(request_finished_at=end, processing_finished_at=end,
                    status="completed" if succeeded is True else "failed" if succeeded is False else "unknown",
                    error_code=None if succeeded is True else "processing_failed" if succeeded is False else "status_unavailable")
        start = data.get("processing_started_at")
        data["processing_ms"] = max(0, end-start) if start else None
    _update(root, jid, change)


def sample_devices():
    """機器全体の瞬時使用率。プロセス単位のGPU時間ではない。"""
    result = subprocess.run([
        "nvidia-smi", "--query-gpu=index,utilization.gpu,memory.used,memory.total",
        "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=1)
    if result.returncode:
        return []
    devices = []
    for line in result.stdout.splitlines():
        values = [value.strip() for value in line.split(",")]
        if len(values) != 4:
            continue
        try:
            index, utilization, used, total = map(float, values)
            if not all(math.isfinite(value) for value in [index, utilization, used, total]) or not 0 <= utilization <= 100 or total <= 0:
                continue
            devices.append({"index": int(index), "utilization_pct": utilization,
                            "memory_used_mib": used, "memory_total_mib": total})
        except ValueError:
            continue
    return devices


def record_gpu_function(function):
    """spaces.GPUの内側に置く。関数の経過時間と機器サンプルを取得する。"""
    signature = inspect.signature(function)
    @functools.wraps(function)
    def wrapped(*args, **kwargs):
        bound = signature.bind(*args, **kwargs)
        jid = bound.arguments.get("jid")
        root = function.__globals__.get("TEMP_DIR")
        if root is None or not re.fullmatch(r"[a-f0-9]{32}", str(jid)):
            return function(*args, **kwargs)
        started = time.perf_counter()
        stop = threading.Event()
        samples = []
        def sample():
            while not stop.is_set():
                try:
                    devices = sample_devices()
                    # 複数GPUは処理対象機器を特定できないため平均へ混ぜない。
                    if len(devices) == 1:
                        samples.append(devices[0]["utilization_pct"])
                except Exception:
                    return
                stop.wait(2)
        thread = threading.Thread(target=sample, daemon=True)
        try:
            thread.start()
        except Exception:
            pass
        try:
            return function(*args, **kwargs)
        finally:
            elapsed = round((time.perf_counter()-started)*1000)
            stop.set()
            def change(data):
                data["gpu_function_ms"] = (data.get("gpu_function_ms") or 0) + elapsed
                # 壁時計時間をGPU稼働時間に変換しない。
                data["gpu_seconds"] = None
                captured = list(samples)
                count = len(captured)
                old_count = data.get("gpu_sample_count", 0)
                old_value = data.get("device_utilization_pct") or 0
                data["gpu_sample_count"] = old_count + count
                data["device_utilization_pct"] = (old_value*old_count+sum(captured))/(old_count+count) if old_count+count else None
            _update(root, jid, change)
    return wrapped