anki_jam.py

from typing import TYPE_CHECKING, Any, Optional

import aqt
import aqt.gui_hooks
import aqt.sound
from anki.cards import Card
from anki.notes import Note
from aqt.utils import tooltip
from PyQt6.QtGui import QAction
from PyQt6.QtWidgets import QCheckBox, QHBoxLayout, QLabel, QMenu, QPushButton, QStatusBar, QWidget

from .riff_importer import RiffImporter
from .riffs_constants import (
    JAM_DECK_NAME,
    JAM_NOTE_FIELDS,
    JAM_NOTETYPE_NAME,
    JAM_PITCH_FIELD,
    JAM_SPEED_FIELD,
    JAM_TEMPLATE_NAME,
)
from .utils import cast_assert

if TYPE_CHECKING:
    from aqt.browser.browser import Browser


def _float_field_or_default(note: Note, field_name: str, default: float) -> float:
    try:
        raw_value = note[field_name]
    except KeyError:
        return default

    stripped = raw_value.strip()
    if not stripped:
        return default

    try:
        return float(stripped)
    except ValueError:
        print(f"Could not parse field '{field_name}' value '{raw_value}' as float.")
        return default


class AudioControls:
    _instance: Optional['AudioControls'] = None

    def __init__(self) -> None:
        self.audio_speed = 1.0
        self.pitch_shift = 0.0  # semitones
        self.show_notif = True
        self.loop_enabled = True
        self.audio_control_widget: Optional['AudioControlWidget'] = None

        self._register_hooks()
        self._register_menu_actions()
        self.ensure_riffs_resources()

    @classmethod
    def instance(cls) -> 'AudioControls':
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    def _register_hooks(self) -> None:
        aqt.gui_hooks.av_player_did_begin_playing.append(self.did_begin_playing)
        aqt.gui_hooks.profile_did_open.append(self.ensure_riffs_resources)
        aqt.gui_hooks.reviewer_did_init.append(self.setup_audio_controls)
        aqt.gui_hooks.card_will_show.append(self.on_card_will_show)
        aqt.gui_hooks.reviewer_will_end.append(self.hide_audio_controls)
        aqt.gui_hooks.browser_will_show_context_menu.append(self.on_browser_will_show_context_menu)

    def _register_menu_actions(self) -> None:
        menu = cast_assert(aqt.mw.form.menuTools.addMenu('Anki Jam'), QMenu)

        action = QAction('Decrease audio speed', aqt.mw)
        action.setShortcut('[')
        action.triggered.connect(self.decrease_audio_speed)
        menu.addAction(action)

        action = QAction('Increase audio speed', aqt.mw)
        action.setShortcut(']')
        action.triggered.connect(self.increase_audio_speed)
        menu.addAction(action)

        action = QAction('Stop audio', aqt.mw)
        action.setShortcut('8')
        action.triggered.connect(self.stop_audio)
        menu.addAction(action)

        action = QAction('Import Loops', aqt.mw)
        action.triggered.connect(self.import_riffs)
        menu.addAction(action)

    def on_browser_will_show_context_menu(self, browser: 'Browser', menu: QMenu) -> None:
        action = QAction('Reset riff speed', browser)
        action.triggered.connect(
            lambda _checked=False, b=browser: self.reset_speed_for_selection(b)
        )
        menu.addAction(action)

    def ensure_riffs_resources(self) -> None:
        self.ensure_riffs_deck_exists()
        self.ensure_riffs_notetype_exists()

    def import_riffs(self) -> None:
        if aqt.mw is None:
            print("Can't import loops, main window is not available")
            return

        if aqt.mw.col is None:
            print("Can't import loops, collection is not available")
            return

        self.ensure_riffs_resources()
        importer = RiffImporter(aqt.mw)
        importer.prompt_and_import()

    def ensure_riffs_deck_exists(self) -> None:
        if aqt.mw is None:
            print("Can't ensure Anki Jam deck exists, main window is not available")
            return

        if aqt.mw.col is None:
            print("Can't ensure Anki Jam deck exists, collection is not available")
            return

        aqt.mw.col.decks.add_normal_deck_with_name(JAM_DECK_NAME)

    def ensure_riffs_notetype_exists(self) -> None:
        mw = aqt.mw
        if mw is None:
            print("Can't ensure Anki Jam note type exists, main window is not available")
            return

        collection = mw.col
        if collection is None:
            print("Can't ensure Anki Jam note type exists, collection is not available")
            return

        models = collection.models
        notetype = models.by_name(JAM_NOTETYPE_NAME)
        if notetype:
            missing_fields = [
                field for field in JAM_NOTE_FIELDS if field not in models.field_map(notetype)
            ]
            if missing_fields:
                for field in missing_fields:
                    models.add_field(notetype, models.new_field(field))
                models.update_dict(notetype)

            if not notetype.get('tmpls'):
                template = models.new_template(JAM_TEMPLATE_NAME)
                template['qfmt'] = '{{Artist}} - {{Title}}<br>{{Name}}<br>{{Audio}}'
                template['afmt'] = '{{FrontSide}}'
                models.add_template(notetype, template)
                models.update_dict(notetype)
            else:
                notetype['tmpls'][0]['name'] = JAM_TEMPLATE_NAME
                notetype['tmpls'][0]['qfmt'] = '{{Artist}} - {{Title}}<br>{{Name}}<br>{{Audio}}'
                notetype['tmpls'][0]['afmt'] = '{{FrontSide}}'
                models.update_dict(notetype)
            return

        notetype = models.new(JAM_NOTETYPE_NAME)
        for field in JAM_NOTE_FIELDS:
            models.add_field(notetype, models.new_field(field))

        template = models.new_template(JAM_TEMPLATE_NAME)
        template['qfmt'] = '{{Artist}} - {{Title}}<br>{{Name}}<br>{{Audio}}'
        template['afmt'] = '{{FrontSide}}'
        models.add_template(notetype, template)
        notetype['sortf'] = 0

        models.add(notetype)

    def reset_speed_for_selection(self, browser: 'Browser') -> None:
        mw = aqt.mw
        if mw is None:
            print("Can't reset riff speeds, main window is not available")
            return

        collection = mw.col
        if collection is None:
            print("Can't reset riff speeds, collection is not available")
            return

        note_ids = browser.selected_notes()
        if not note_ids:
            self.notify('No notes selected to reset speed')
            return

        updated = 0
        skipped = 0
        for note_id in note_ids:
            note = collection.get_note(note_id)
            notetype = note.note_type()
            if not notetype or notetype.get('name') != JAM_NOTETYPE_NAME:
                skipped += 1
                continue

            try:
                note[JAM_SPEED_FIELD] = '1.000'
            except KeyError:
                skipped += 1
                continue

            try:
                collection.update_note(note)
            except AttributeError:
                note.flush()

            updated += 1

        if updated:
            try:
                browser.model.reset()
            except AttributeError:
                browser.mw.reset()
            self.notify(f'Reset speed to 1.0x for {updated} note(s)')
        elif skipped:
            self.notify('No Anki Jam notes in selection to reset speed')

    def set_speed(self, player: Optional[aqt.sound.MpvManager], speed: float) -> None:
        if player is None:
            print("Couldn't set speed, player is None")
            return

        self.audio_speed = speed
        player_ = cast_assert(player, aqt.sound.MpvManager)
        player_.set_property('af', self._build_audio_filter_value())
        player_.set_property('speed', self.audio_speed)

    def apply_audio_speed(self, player: Any, speed: float) -> None:
        self.set_speed(player, speed)

    def apply_pitch_shift(self, player: Any, pitch_shift: float) -> None:
        if player is None:
            print("Couldn't set pitch shift, player is None")
            return

        self.pitch_shift = pitch_shift
        player_ = cast_assert(player, aqt.sound.MpvManager)
        player_.set_property('af', self._build_audio_filter_value())
        player_.set_property('speed', self.audio_speed)

    def apply_loop_setting(self, player: Any) -> None:
        if player is None:
            return

        if isinstance(player, aqt.sound.MpvManager):
            player.set_property('loop', 'inf' if self.loop_enabled else 'no')  # type: ignore[no-untyped-call]

    def toggle_loop(self, enabled: bool) -> None:
        self.loop_enabled = enabled
        self.apply_loop_setting(aqt.sound.av_player.current_player)
        self.notify('Loop enabled' if enabled else 'Loop disabled')
        self.try_update_displays()

    def increase_audio_speed(self) -> None:
        self.audio_speed = min(4.0, self.audio_speed + 0.1)
        self.notify('Increase audio speed by 0.1. Current Speed: ' + f'{self.audio_speed:.1f}')

        self.apply_audio_speed(aqt.sound.av_player.current_player, self.audio_speed)
        self.try_update_displays()

    def decrease_audio_speed(self) -> None:
        self.audio_speed = max(0.1, self.audio_speed - 0.1)
        self.notify('Decrease audio speed by 0.1. Current Speed: ' + f'{self.audio_speed: .1f}')

        self.apply_audio_speed(aqt.sound.av_player.current_player, self.audio_speed)
        self.try_update_displays()

    def stop_audio(self) -> None:
        self.notify('Stop audio.')

        player = aqt.sound.av_player.current_player
        if player is None:
            print("Can't stop audio, player is None")
            return

        cast_assert(player, aqt.sound.MpvManager).stop()

    def did_begin_playing(self, player: Any, _tag: Any) -> None:
        print('BEGIN PLAYING')
        player.set_property('audio-delay', '0')
        player.set_property('gapless-audio', 'yes')
        self.apply_pitch_shift(player, self.pitch_shift)
        self.apply_audio_speed(player, self.audio_speed)
        self.apply_loop_setting(player)

    def on_card_will_show(self, text: str, card: Card, _kind: str) -> str:
        self.apply_card_audio_settings(card)
        return text

    def apply_card_audio_settings(self, card: Card) -> None:
        note = card.note()
        notetype = note.note_type()
        if not notetype or notetype.get('name') != JAM_NOTETYPE_NAME:
            return

        pitch_value = _float_field_or_default(note, JAM_PITCH_FIELD, 0.0)
        speed_value = _float_field_or_default(note, JAM_SPEED_FIELD, 1.0)

        self.pitch_shift = pitch_value
        self.audio_speed = speed_value
        player = aqt.sound.av_player.current_player
        if player:
            self.apply_pitch_shift(player, self.pitch_shift)
            self.apply_audio_speed(player, self.audio_speed)
        self.try_update_displays()

    def _build_audio_filter_value(self) -> str:
        pitch_scale = 2 ** (self.pitch_shift / 12)
        return f'rubberband=window=long:formant=preserved:pitch-scale={pitch_scale:.6f}'

    def setup_audio_controls(self, _reviewer: Any) -> None:
        if self.audio_control_widget is None:
            self.audio_control_widget = AudioControlWidget(self, aqt.mw)
            status_bar = cast_assert(aqt.mw.statusBar(), QStatusBar)
            status_bar.addPermanentWidget(self.audio_control_widget)

        self.audio_control_widget.show()
        self.try_update_displays()

    def hide_audio_controls(self) -> None:
        if self.audio_control_widget:
            self.audio_control_widget.hide()

    def notify(self, message: str) -> None:
        if self.show_notif:
            tooltip(message)

    def try_update_displays(self) -> None:
        if self.audio_control_widget:
            self.audio_control_widget.update_displays()


class AudioControlWidget(QWidget):
    def __init__(self, controller: AudioControls, parent: Optional[QWidget] = None) -> None:
        super().__init__(parent)
        self.controller = controller
        self.setup_ui()

    def setup_ui(self) -> None:
        layout = QHBoxLayout()
        layout.setContentsMargins(5, 2, 5, 2)

        self.speed_label = QLabel('Speed: 1.0x')
        self.speed_dec_btn = QPushButton('-')
        self.speed_inc_btn = QPushButton('+')
        self.speed_dec_btn.setMaximumWidth(30)
        self.speed_inc_btn.setMaximumWidth(30)

        self.pitch_label = QLabel('Pitch: 0.0')
        self.pitch_dec_btn = QPushButton('-')
        self.pitch_inc_btn = QPushButton('+')
        self.pitch_dec_btn.setMaximumWidth(30)
        self.pitch_inc_btn.setMaximumWidth(30)

        self.loop_checkbox = QCheckBox('Loop')
        self.loop_checkbox.setChecked(self.controller.loop_enabled)
        self.loop_checkbox.toggled.connect(self.controller.toggle_loop)

        self.speed_dec_btn.clicked.connect(self.controller.decrease_audio_speed)
        self.speed_inc_btn.clicked.connect(self.controller.increase_audio_speed)

        layout.addWidget(self.speed_label)
        layout.addWidget(self.speed_dec_btn)
        layout.addWidget(self.speed_inc_btn)
        layout.addSpacing(20)
        layout.addWidget(self.pitch_label)
        layout.addWidget(self.pitch_dec_btn)
        layout.addWidget(self.pitch_inc_btn)
        layout.addSpacing(20)
        layout.addWidget(self.loop_checkbox)
        layout.addStretch()

        self.setLayout(layout)

    def update_displays(self) -> None:
        self.speed_label.setText(f'Speed: {self.controller.audio_speed:.1f}x')
        if abs(self.controller.pitch_shift) < 0.0001:
            self.pitch_label.setText('Pitch: 0.0')
        else:
            self.pitch_label.setText(f'Pitch: {self.controller.pitch_shift:+.1f}')
        self.loop_checkbox.setChecked(self.controller.loop_enabled)


audio_controls = AudioControls.instance()