#!/usr/bin/env python3
"""Pinned PeerTube runner 0.2.0 FFmpeg adapter (native Windows).

Port of infra/peertube-workers/ffmpeg-qsv.py from the validated Linux QSV worker.
Same contract: rewrite known runner command shapes into explicit QSV decode ->
hardware frames -> scale_qsv -> h264_qsv; never execute a CPU video encoder;
fail closed (exit 78) on anything unreviewed.

Windows deltas from the Linux adapter:
  - real ffmpeg/ffprobe are sibling .exe files of this shim (the runner resolves
    FFMPEG_PATH with fs.exists, then spawns it without a shell, so the shim
    MUST be a real .exe: a .cmd/.bat wrapper is rejected by Node >= 18.20/20.12
    with EINVAL, and fluent-ffmpeg passes no `shell` option).
  - QSV device init is `-init_hw_device qsv=qs` (D3D11 child device by default;
    the libmfx/oneVPL dispatcher DLLs must sit next to ffmpeg.exe).
  - "hardware really did it" telemetry is the PDH GPU Engine counter
    \\GPU Engine(pid_<pid>_*)\\Utilization Percentage over engtype_VideoEncode
    instances, sampled via typeperf (i915 DRM fdinfo does not exist on Windows).
    If the counters are readable and stay exactly 0 for the whole run, the run
    is recorded as a GPU failure even on exit 0.
  - capability queries (-version/-encoders/...) are passed through untouched.
"""
import json
import os
import re
import subprocess
import sys
import threading
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
REAL = Path(os.environ.get('PT_FFMPEG_EXE', str(HERE / 'ffmpeg.exe')))
PROBE = Path(os.environ.get('PT_FFPROBE_EXE', str(HERE / 'ffprobe.exe')))
# Optional child-device spec for the QSV device (e.g. 'd3d11va', 'd3d11va:1').
# Empty means the default adapter. Verify on first hardware contact.
DEVICE = os.environ.get('PT_QSV_DEVICE', '').strip()
STATE = Path(os.environ.get('PT_STATE_DIR', r'C:\pt-worker\state'))

PASSTHROUGH = (['-version'], ['-formats'], ['-encoders'], ['-filters'],
               ['-codecs'], ['-hwaccels'])


def event(kind, **fields):
    STATE.mkdir(parents=True, exist_ok=True)
    line = json.dumps({'time': time.time(), 'kind': kind, **fields}) + '\n'
    try:
        with open(STATE / 'events.jsonl', 'a', encoding='utf8') as handle:
            handle.write(line)
    except OSError:
        pass  # telemetry loss must never fail a job


def probe(path):
    p = subprocess.run([str(PROBE), '-v', 'error', '-show_streams', '-show_format',
                        '-of', 'json', str(path)], capture_output=True, text=True, timeout=30)
    if p.returncode:
        raise ValueError('input probe failed')
    return json.loads(p.stdout)


def csv_cells(line):
    return [cell.strip('"') for cell in line.split('","')]


def as_float(text):
    try:
        return float(text)
    except (TypeError, ValueError):
        pass
    try:  # locale decimal-comma fallback
        return float(text.replace('.', '').replace(',', '.'))
    except (TypeError, ValueError):
        return None


def video_encode_busy(pid):
    """Sum of PDH \\GPU Engine(pid_<pid>_*)\\Utilization Percentage over
    engtype_VideoEncode instances, one 1-second typeperf sample.

    Returns (busy_or_None, readable). Counter names and decimal points are
    locale-sensitive; unparseable samples count as not-readable, never as zero.
    """
    counter = rf'\GPU Engine(pid_{pid}_*)\Utilization Percentage'
    try:
        p = subprocess.run(['typeperf', counter, '-sc', '1'], capture_output=True,
                           text=True, timeout=15)
    except (OSError, subprocess.SubprocessError):
        return None, False
    if p.returncode:
        return None, False
    rows = [line for line in (p.stdout or '').splitlines() if line.startswith('"')]
    if len(rows) < 2:
        return None, False
    headers = list(csv_cells(rows[0]))
    values = list(csv_cells(rows[1]))
    total, seen = 0.0, False
    for index, header in enumerate(headers):
        if 'engtype_VideoEncode' not in header.lower() or index >= len(values):
            continue
        number = as_float(values[index])
        if number is None:
            return None, False
        seen = True
        total += number
    return (total if seen else None), seen


def rewrite(args, input_probe=None):
    # fluent-ffmpeg queries capabilities before building commands. These do no encode.
    if list(args) in [list(p) for p in PASSTHROUGH]:
        return args, None
    pairs = list(zip(args, args[1:]))
    video_options = [(k, v) for k, v in pairs if k in ('-vcodec', '-c:v', '-codec:v')]
    if len(video_options) > 1:
        raise ValueError('multiple video codec options refused')
    if '-vn' in args and not video_options:
        return args, None  # audio-only VOD renditions have no video encode
    if not video_options or video_options[0][1] != 'libx264':
        raise ValueError('unknown runner command or video encoder refused')
    if '-filter_complex' in args or '-hwaccel' in args or '-init_hw_device' in args:
        raise ValueError('unreviewed filter/device override refused')
    inputs = [v for k, v in pairs if k == '-i']
    if not inputs:
        raise ValueError('missing video input')
    data = input_probe if input_probe is not None else probe(inputs[0])
    stream = next((s for s in data['streams'] if s['codec_type'] == 'video'), None)
    if stream is None:
        raise ValueError('missing video stream')
    codec = stream['codec_name']
    still = codec in ('png', 'mjpeg', 'webp') and '-loop' in args
    decoder = {'h264': 'h264_qsv', 'hevc': 'hevc_qsv', 'vp9': 'vp9_qsv', 'av1': 'av1_qsv'}.get(codec)
    if not decoder and not still:
        raise ValueError('source codec has no reviewed QSV decode path')
    if stream.get('pix_fmt') not in ('yuv420p', 'yuv420p10le', 'nv12', 'p010le') and not still:
        raise ValueError('unsupported pixel format; CPU decode fallback refused')
    if stream.get('tags', {}).get('rotate') or any(s.get('rotation', 0) for s in stream.get('side_data_list', [])):
        raise ValueError('rotated source requires reviewed GPU transpose path')
    width, height = stream['width'], stream['height']
    duration = float(data.get('format', {}).get('duration', 0) or 0)
    w, h = width // 2 * 2, height // 2 * 2
    output = ['-init_hw_device', f'qsv=qs:{DEVICE}' if DEVICE else 'qsv=qs',
              '-filter_hw_device', 'qs']
    i, first_input = 0, True
    while i < len(args):
        key = args[i]
        val = args[i + 1] if i + 1 < len(args) else None
        if key == '-i' and first_input:
            if not still:
                output += ['-hwaccel', 'qsv', '-hwaccel_output_format', 'qsv', '-c:v', decoder, '-noautorotate']
            output += [key, val]
            first_input = False
            i += 2
        elif key in ('-vcodec', '-c:v', '-codec:v'):
            output += [key, 'h264_qsv']
            i += 2
        elif key in ('-pix_fmt', '-preset', '-preset:v', '-b_strategy', '-bf', '-tune', '-crf'):
            i += 2
        elif key in ('-vf', '-filter:v'):
            if val == 'scale=trunc(iw/2)*2:trunc(ih/2)*2' and still:
                w, h = width // 2 * 2, height // 2 * 2
            else:
                match = re.fullmatch(r'scale=w=(-2|\d+):h=(-2|\d+)', val or '')
                if not match:
                    raise ValueError('unreviewed software filter refused')
                w, h = map(int, match.groups())
                if w == -2 and h > 0:
                    w = max(2, round(width * h / height / 2) * 2)
                elif h == -2 and w > 0:
                    h = max(2, round(height * w / width / 2) * 2)
                if w <= 0 or h <= 0 or w % 2 or h % 2:
                    raise ValueError('invalid output dimensions')
            prefix = 'format=nv12,hwupload=extra_hw_frames=64,' if still else ''
            output += [key, prefix + f'scale_qsv=w={w}:h={h}:format=nv12']
            i += 2
        else:
            output.append(key)
            i += 1
    if '-vf' not in args and '-filter:v' not in args:
        prefix = 'format=nv12,hwupload=extra_hw_frames=64,' if still else ''
        output[-1:-1] = ['-vf', prefix + f'scale_qsv=w={width // 2 * 2}:h={height // 2 * 2}:format=nv12']
    # Keep PeerTube's rate cap and GOP; QSV has no x264 CRF mode.
    rate = next((v for k, v in pairs if k in ('-maxrate:v', '-maxrate')), None)
    extra = ['-preset', 'medium', '-bf', '2']
    if rate and not any(k in ('-b:v', '-vb') for k, _ in pairs):
        extra += ['-b:v', rate]
    output[-1:-1] = extra
    job_uuid = None
    try:
        job_uuid = json.loads((STATE / 'inputs' / (Path(inputs[0]).name + '.json')).read_text(encoding='utf8'))['uuid']
    except (OSError, ValueError, KeyError):
        pass
    return output, {'job_uuid': job_uuid, 'source_codec': codec, 'target_codec': 'h264', 'encoder': 'h264_qsv',
                    'source_width': width, 'source_height': height, 'duration': duration,
                    'target_width': w, 'target_height': h,
                    'static_image_upload': still}


def main():
    try:
        if not REAL.is_file():
            raise ValueError('real ffmpeg.exe not found next to shim')
        args, info = rewrite(sys.argv[1:])
        if info is None:
            os.execv(str(REAL), [str(REAL), *args])
        started = time.monotonic()
        p = subprocess.Popen([str(REAL), *args], stderr=subprocess.PIPE, text=True,
                             encoding='utf8', errors='replace')
        progress = {}
        def stderr_progress():
            for line in p.stderr:
                sys.stderr.write(line)
                sys.stderr.flush()
                for key, value in re.findall(r'(frame|fps|speed)=\s*([\d.]+)', line):
                    try:
                        progress[key] = float(value)
                    except ValueError:
                        pass
        reader = threading.Thread(target=stderr_progress, daemon=True)
        reader.start()
        maximum, samples, readable = None, 0, False
        while p.poll() is None:
            busy, ok = video_encode_busy(p.pid)
            readable |= ok
            if ok:
                samples += 1
                if busy is not None:
                    maximum = busy if maximum is None else max(maximum, busy)
            time.sleep(1.0)
        elapsed = time.monotonic() - started
        reader.join(timeout=2)
        event('encode', **info, seconds=elapsed, returncode=p.returncode,
              engine_busy_max=maximum, engine_samples=samples,
              engine_counters_readable=readable,
              ffmpeg_fps=progress.get('fps'), frames=progress.get('frame'),
              ffmpeg_speed=progress.get('speed'),
              realtime_factor=info['duration'] / elapsed if info['duration'] else None)
        if p.returncode:
            event('gpu_failure', reason='qsv_transcode_failed')
        elif readable and samples >= 3 and maximum == 0:
            # Exit 0 with a readable VideoEncode counter pinned at 0 is NOT
            # proof of hardware encoding; fail closed like the Linux fdinfo gate.
            event('gpu_failure', reason='no_video_encode_utilization')
        sys.exit(p.returncode)
    except (ValueError, OSError, subprocess.SubprocessError, KeyError) as e:
        # Messages are fixed strings; never expose source URLs or runner credentials.
        event('gpu_failure', reason=type(e).__name__)
        print('PT-WORKER: QSV adapter refused command; worker must drain', file=sys.stderr)
        sys.exit(78)


if __name__ == '__main__':
    main()