#!/usr/bin/python3
# SPDX-License-Identifier: MIT
"""PianoRT G3 tuning: CPU profile, IRQ placement, synthesis thread placement
and audio DSP power. Settings come from /etc/pianort/tuning.conf.

  g3-tuning plan          show what would be applied, change nothing
  g3-tuning profile       SMT, CPU frequency and package power
  g3-tuning irq           IRQ affinity and IRQ thread priorities
  g3-tuning watch         keep Pianoteq's audio threads on their CPUs
  g3-tuning audio-power   keep the SOF audio DSP from runtime-suspending
"""
import json
import os
import pathlib
import re
import shlex
import subprocess
import sys
import time

CPU = pathlib.Path('/sys/devices/system/cpu')
CONF = pathlib.Path('/etc/pianort/tuning.conf')
PLAYER = pathlib.Path('/etc/pianort/player.conf')
RAPL = pathlib.Path('/sys/class/powercap/intel-rapl:0')
DSP_POWER = pathlib.Path('/sys/bus/pci/devices/0000:00:1f.3/power/control')

def read(path):
    return pathlib.Path(path).read_text().strip()

def write(path, value):
    """Write a sysfs value if the file exists; report failures without stopping."""
    path = pathlib.Path(path)
    if not path.exists():
        return False
    try:
        path.write_text(str(value))
        return True
    except OSError as error:
        print('PianoRT: could not write ' + str(path) + ': ' + str(error), file=sys.stderr, flush=True)
        return False

def config(path):
    values = {}
    try:
        text = path.read_text()
    except FileNotFoundError:
        return values
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith('#') or '=' not in line:
            continue
        key, _, value = line.partition('=')
        values[key.strip()] = ' '.join(shlex.split(value)) if value.strip() else ''
    return values

def cpulist(text):
    cpus = set()
    for part in text.split(','):
        part = part.strip()
        if part:
            low, _, high = part.partition('-')
            cpus.update(range(int(low), int(high or low) + 1))
    return cpus

def fmt(cpus):
    return ','.join(str(c) for c in sorted(cpus))

def online():
    return cpulist(read(CPU / 'online'))

def physical_cores():
    cores = set()
    for cpu in CPU.glob('cpu[0-9]*'):
        core = cpu / 'topology/core_id'
        if core.exists():
            package = cpu / 'topology/physical_package_id'
            cores.add((read(package) if package.exists() else '0', read(core)))
    return len(cores)

def layout():
    """Return (housekeeping, audio) CPU lists.

    Uses the isolcpus= list on the kernel command line, as written by the
    PianoRT firmware flasher. The sysfs views can't be used instead: with only
    the managed_irq flag, /sys/.../cpu/isolated stays empty, and nohz_full
    always drops the boot CPU. Without isolcpus, the last CPU does
    housekeeping and the rest do audio, matching the reference unit.
    """
    up = online()
    match = re.search(r'(?:^|\s)isolcpus=(\S+)', read('/proc/cmdline'))
    if match:
        # Drop flags such as managed_irq or domain; keep the CPU list.
        cpus = ','.join(p for p in match.group(1).split(',') if p and p[0].isdigit())
        isolated = cpulist(cpus) & up
        if isolated and up - isolated:
            return sorted(up - isolated), sorted(isolated)
    last = max(up)
    return [last], sorted(up - {last}) or [last]

def irqs(pattern):
    found = []
    for line in pathlib.Path('/proc/interrupts').read_text().splitlines()[1:]:
        number, _, rest = line.partition(':')
        if number.strip().isdigit() and re.search(pattern, rest, re.IGNORECASE):
            found.append(number.strip())
    return found

def threads(pattern):
    out = subprocess.run(['ps', '-eLo', 'tid=,comm='], capture_output=True, text=True, check=True).stdout
    return [line.split(None, 1)[0] for line in out.splitlines()
            if len(line.split(None, 1)) == 2 and re.search(pattern, line.split(None, 1)[1])]

def player_executable():
    target = config(PLAYER).get('PIANORT_EXECUTABLE') or '/opt/pianoteq/current/Pianoteq'
    return os.path.realpath(target)

# --- plan -----------------------------------------------------------------------

def plan(conf):
    housekeeping, audio = layout()
    print(json.dumps({
        'physical_cores': physical_cores(),
        'online': fmt(online()),
        'housekeeping_cpus': fmt(housekeeping),
        'audio_cpus': fmt(audio),
        'smt_off': conf.get('PIANORT_SMT_OFF', 'yes') == 'yes' and physical_cores() >= 4,
        'cpu_khz': conf.get('PIANORT_CPU_KHZ', ''),
        'package_power_w': conf.get('PIANORT_PACKAGE_POWER_W', ''),
        'wifi_usb_irqs_to': fmt(housekeeping), 'wifi_usb_irqs': irqs(r'iwlwifi|xhci'),
        'audio_dsp_irqs_to': fmt(audio), 'audio_dsp_irqs': irqs(r'AudioDSP'),
        'thread_placement': physical_cores() >= 4,
        'pianoteq': player_executable(),
        'pin_processes': conf.get('PIANORT_PIN_PROCESSES', ''),
    }, indent=2))

# --- profile ----------------------------------------------------------------------

def profile(conf):
    smt = CPU / 'smt/control'
    if conf.get('PIANORT_SMT_OFF', 'yes') == 'yes' and physical_cores() >= 4 and smt.exists():
        if read(smt) == 'on':
            write(smt, 'off')
    up = online()
    want = int(conf.get('PIANORT_CPU_KHZ') or 0)
    applied_khz = None
    for policy in sorted((CPU / 'cpufreq').glob('policy*')):
        if int(policy.name[6:]) not in up:
            continue
        write(policy / 'scaling_governor', 'performance')
        write(policy / 'energy_performance_preference', 'performance')
        if want:
            khz = min(want, int(read(policy / 'cpuinfo_max_freq')))
            khz = max(khz, int(read(policy / 'cpuinfo_min_freq')))
            if int(read(policy / 'scaling_min_freq')) > khz:
                write(policy / 'scaling_min_freq', khz)
            write(policy / 'scaling_max_freq', khz)
            write(policy / 'scaling_min_freq', khz)
            applied_khz = khz
    power_uw = None
    watts = conf.get('PIANORT_PACKAGE_POWER_W', '')
    if watts:
        limit = RAPL / 'constraint_0_power_limit_uw'
        for _ in range(100):
            if limit.exists():
                break
            time.sleep(0.1)
        # Only the sustained package limit (PL1) changes. PL2, time windows and
        # platform limits stay as the firmware set them.
        if limit.exists() and read(RAPL / 'name') == 'package-0':
            value = str(int(float(watts) * 1000000))
            if write(limit, value) and read(limit) == value:
                power_uw = int(value)
            else:
                print('PianoRT: package power limit readback failed; firmware policy kept', file=sys.stderr, flush=True)
        else:
            print('PianoRT: package powercap unavailable; firmware power policy kept', file=sys.stderr, flush=True)
    print(json.dumps({'smt': read(smt) if smt.exists() else 'n/a', 'online': fmt(online()),
                      'frequency_khz': applied_khz, 'package_sustained_power_uw': power_uw}), flush=True)

# --- irq ------------------------------------------------------------------------------

def irq(conf):
    housekeeping, audio = layout()
    # Wi-Fi and USB interrupts on the housekeeping CPU, audio DSP on the audio CPUs.
    for number in irqs(r'iwlwifi|xhci'):
        write('/proc/irq/' + number + '/smp_affinity_list', fmt(housekeeping))
    for number in irqs(r'AudioDSP'):
        write('/proc/irq/' + number + '/smp_affinity_list', fmt(audio))
    # Wi-Fi IRQ threads to SCHED_IDLE so they never preempt anything.
    for tid in threads(r'iwlwifi'):
        subprocess.run(['chrt', '-i', '-p', '0', tid], capture_output=True)
    # Audio DSP IRQ above every Pianoteq thread (its data loop runs at FIFO 94).
    for tid in threads(r'AudioDSP'):
        subprocess.run(['chrt', '-f', '-p', '99', tid], capture_output=True)
    # NVMe queues stay on the kernel's managed spreading; their IRQ threads go
    # above stock priority 50 but below Pianoteq's realtime stack (64/90/94).
    for tid in threads(r'nvme[0-9]+q'):
        subprocess.run(['chrt', '-f', '-p', '55', tid], capture_output=True)
    print(json.dumps({'housekeeping_cpus': fmt(housekeeping), 'audio_cpus': fmt(audio)}), flush=True)

# --- watch --------------------------------------------------------------------------

def is_player(pid, target):
    try:
        return pathlib.Path('/proc', str(pid)).stat().st_uid != 0 and os.readlink('/proc/' + str(pid) + '/exe') == target
    except OSError:
        return False

def place(pid, housekeeping, audio):
    changes = []
    for task in pathlib.Path('/proc', str(pid), 'task').iterdir():
        try:
            tid = int(task.name)
            name = read(task / 'comm')
            if name == 'data-loop.0':
                cpu = housekeeping[0]
            elif re.fullmatch(r'fasthp-[0-9]+', name):
                cpu = audio[(int(name.split('-')[1]) - 1) % len(audio)]
            else:
                continue
            if os.sched_getaffinity(tid) != {cpu}:
                os.sched_setaffinity(tid, {cpu})
                changes.append({'tid': tid, 'name': name, 'cpu': cpu})
        except (OSError, ProcessLookupError, ValueError):
            continue
    return changes

def pin_processes(names, housekeeping):
    for proc in pathlib.Path('/proc').iterdir():
        if not proc.name.isdigit():
            continue
        try:
            if read(proc / 'comm') in names and os.sched_getaffinity(int(proc.name)) != set(housekeeping):
                os.sched_setaffinity(int(proc.name), set(housekeeping))
        except (OSError, ProcessLookupError):
            continue

def watch(conf):
    if physical_cores() < 4:
        print('PianoRT: thread placement is tested on 4-core G3 models only; not placing', flush=True)
        while True:
            time.sleep(3600)
    housekeeping, audio = layout()
    os.sched_setaffinity(0, set(housekeeping))
    target = player_executable()
    names = set(conf.get('PIANORT_PIN_PROCESSES', '').split())
    pids = []
    next_discovery = 0
    while True:
        if time.monotonic() >= next_discovery or not pids or not all(is_player(p, target) for p in pids):
            pids = [int(p.name) for p in pathlib.Path('/proc').iterdir() if p.name.isdigit() and is_player(p.name, target)]
            if names:
                pin_processes(names, housekeeping)
            next_discovery = time.monotonic() + 2
        for pid in pids:
            try:
                changes = place(pid, housekeeping, audio)
                if changes:
                    print(json.dumps({'time': time.time(), 'pid': pid, 'changes': changes}), flush=True)
            except FileNotFoundError:
                pass
        time.sleep(0.5)

# --- audio-power ----------------------------------------------------------------------

def audio_power(conf):
    # Wait for the ALSA card itself: the SOF driver's async probe resets
    # power/control to "auto" shortly after the device appears.
    for _ in range(60):
        try:
            if 'sofrt5682' in pathlib.Path('/proc/asound/cards').read_text():
                break
        except FileNotFoundError:
            pass
        time.sleep(0.5)
    write(DSP_POWER, 'on')

def main():
    commands = {'plan': plan, 'profile': profile, 'irq': irq, 'watch': watch, 'audio-power': audio_power}
    if len(sys.argv) != 2 or sys.argv[1] not in commands:
        print(__doc__.strip(), file=sys.stderr)
        raise SystemExit(2)
    commands[sys.argv[1]](config(CONF))

if __name__ == '__main__':
    main()
