riff_importer.py

import hashlib
import re
import subprocess
import sys
import tempfile
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator, List, NamedTuple, NewType, Optional, Sequence, Tuple

import aqt
from anki.decks import DeckId
from anki.models import NotetypeDict
from aqt.utils import getFile, showWarning, tooltip
from PyQt6.QtCore import Qt
from PyQt6.QtWidgets import QProgressDialog

from .riffs_constants import JAM_DECK_NAME, JAM_NOTETYPE_NAME
from .utils import cast_assert

ChunkHash = NewType('ChunkHash', str)


class ImportResult(NamedTuple):
    created: int
    skipped: int


@dataclass
class XscDocument:
    audio_file: 'AudioFile'
    loops: List['Loop']
    tempo_adjustment: float
    eq_settings: Optional[List[float]]


@dataclass
class AudioFile:
    path: Path
    sample_rate: int
    channels: int


@dataclass
class Loop:
    id: int
    enabled: bool
    start_samples: int
    duration_samples: int
    color: str


@dataclass
class PitchVariant:
    shift: float
    alterations: int


@dataclass
class Section:
    name: str
    lines: List['Line']

    def find_line(self, key: str) -> 'Line':
        for line in self.lines:
            if line.key == key:
                return line
        raise ValueError(f'Missing line: {key}')


@dataclass
class Line:
    key: str
    value: str


class RiffImporter:
    def __init__(self, mw: aqt.AnkiQt) -> None:  # type: ignore[name-defined]
        self.mw = mw

    def prompt_and_import(self) -> None:
        xsc_path_str = getFile(
            self.mw,
            'Select Transcribe! session (.xsc)',
            None,
            key='control_audio_last_xsc',
            filter='Transcribe Session (*.xsc)',
        )
        if not xsc_path_str:
            return

        xsc_path = Path(str(xsc_path_str))
        try:
            result = self.import_file(xsc_path)
        except Exception as exc:  # noqa: BLE001
            showWarning(f'Import failed: {exc}')
            return

        tooltip(f'Imported riffs: {result.created} new, {result.skipped} skipped')
        self.mw.reset()

    def import_file(self, xsc_path: Path) -> ImportResult:
        collection = self.mw.col
        if collection is None:
            raise RuntimeError('Collection is not available')

        document = parse_xsc_file(xsc_path)
        if not document.audio_file.path.exists():
            raise RuntimeError(
                f'Audio file not found: {document.audio_file.path} (from {xsc_path.name})'
            )
        if document.audio_file.sample_rate <= 0:
            raise RuntimeError('Invalid sample rate in .xsc file')
        enabled_loops = sorted(
            (loop for loop in document.loops if loop.enabled),
            key=lambda loop: loop.duration_samples,
        )
        if not enabled_loops:
            raise RuntimeError('No enabled loops were found in this .xsc file')

        mpv_path = resolve_mpv_path()
        deck_id = self._deck_id()
        notetype = self._notetype()
        existing_hashes = self._existing_hashes()

        created = 0
        skipped = 0

        with tempfile.TemporaryDirectory(prefix='anki-riff-import-') as temp_dir:
            tmp_root = Path(temp_dir)
            self.mw.progress.start(immediate=True, label='Importing riffs...')
            try:
                for index, loop in enumerate(enabled_loops, start=1):
                    self.mw.progress.update(
                        label=f'Importing loop {index}/{len(enabled_loops)}',
                        value=index,
                        max=len(enabled_loops),
                    )

                    speed = document.tempo_adjustment
                    pitch_variants = build_pitch_shift_sequence()
                    chunk_hashes = [
                        compute_hash(
                            xsc_path.name,
                            loop.start_samples,
                            loop.start_samples + loop.duration_samples,
                            float(variant.shift),
                            speed,
                        )
                        for variant in pitch_variants
                    ]

                    if all(chunk_hash in existing_hashes for chunk_hash in chunk_hashes):
                        skipped += len(chunk_hashes)
                        continue

                    start_sec = loop.start_samples / document.audio_file.sample_rate
                    end_sec = (
                        loop.start_samples + loop.duration_samples
                    ) / document.audio_file.sample_rate
                    duration_sec = end_sec - start_sec

                    if duration_sec <= 0:
                        skipped += len(chunk_hashes)
                        continue

                    audio_filename = build_media_filename(chunk_hashes[0])
                    temp_audio_path = tmp_root / audio_filename

                    with self._transcoding_dialog(
                        f'Transcoding loop {index}/{len(enabled_loops)}'
                    ):
                        export_chunk_with_mpv(
                            mpv_path,
                            document.audio_file.path,
                            start_sec,
                            end_sec,
                            temp_audio_path,
                        )

                    media_data = temp_audio_path.read_bytes()
                    media_name = collection.media.write_data(audio_filename, media_data)

                    for variant, chunk_hash in zip(pitch_variants, chunk_hashes):
                        if chunk_hash in existing_hashes:
                            skipped += 1
                            continue

                        self._add_note(
                            notetype,
                            deck_id,
                            media_name,
                            xsc_path,
                            document,
                            loop,
                            chunk_hash,
                            float(variant.shift),
                            variant.alterations,
                            speed,
                            start_sec,
                            end_sec,
                            duration_sec,
                        )
                        existing_hashes.add(chunk_hash)
                        created += 1
            finally:
                self.mw.progress.finish()

        return ImportResult(created=created, skipped=skipped)

    def _deck_id(self) -> int:
        collection = self.mw.col
        if collection is None:
            raise RuntimeError('Collection is not available')

        deck_id_obj = collection.decks.add_normal_deck_with_name(JAM_DECK_NAME)
        return int(deck_id_obj.id)

    def _notetype(self) -> NotetypeDict:
        collection = self.mw.col
        if collection is None:
            raise RuntimeError('Collection is not available')

        models = collection.models
        notetype = cast_assert(models.by_name(JAM_NOTETYPE_NAME), NotetypeDict)

        return notetype  # type: ignore[no-any-return]

    def _existing_hashes(self) -> set[str]:
        collection = self.mw.col
        if collection is None:
            return set()

        hashes: set[str] = set()
        for note_id in collection.find_notes(f'note:"{JAM_NOTETYPE_NAME}"'):
            note = collection.get_note(note_id)
            value = note['hash'].strip()
            if value:
                hashes.add(value)

        return hashes

    @contextmanager
    def _transcoding_dialog(self, message: str) -> Iterator[None]:
        dialog = QProgressDialog(message, '', 0, 0, self.mw)
        dialog.setWindowModality(Qt.WindowModality.ApplicationModal)
        dialog.setCancelButton(None)
        dialog.setMinimumDuration(0)
        dialog.setAutoClose(False)
        dialog.setAutoReset(False)
        dialog.setWindowTitle('Transcoding riff audio')
        dialog.show()

        self.mw.app.processEvents()

        yield

        dialog.close()

    def _add_note(
        self,
        notetype: Any,
        deck_id: int,
        media_name: str,
        xsc_path: Path,
        document: XscDocument,
        loop: Loop,
        chunk_hash: ChunkHash,
        pitch_shift: float,
        alterations: int,
        speed: float,
        start_sec: float,
        end_sec: float,
        duration_sec: float,
    ) -> None:
        collection = self.mw.col
        if collection is None:
            raise RuntimeError('Collection is not available')

        note = collection.new_note(notetype)
        artist, title_root = split_artist_title(xsc_path.stem)
        riff_name = build_riff_name(xsc_path, loop.id)
        color = text_color_for_shift(pitch_shift, alterations)

        if not artist:
            artist = title_root if title_root else 'Unknown'

        # At some point, make this an option
        speed = 1.0

        note['Artist'] = wrap_with_color(artist, color)
        note['Title'] = wrap_with_color(title_root or riff_name, color)
        note['Name'] = wrap_with_color(riff_name, color)
        note['Audio'] = f'[sound:{media_name}]'
        note['pitch-shift'] = f'{pitch_shift:.3f}'
        note['speed'] = f'{speed:.3f}'
        note['hash'] = chunk_hash
        note['original path'] = str(document.audio_file.path)
        note['original filename'] = document.audio_file.path.name
        note['chunk start'] = format_timestamp(start_sec)
        note['chunk end'] = format_timestamp(end_sec)
        note['duration'] = f'{duration_sec:.3f}'

        song_tag = build_song_tag(document.audio_file.path)
        shift_tag = build_shift_tag(pitch_shift)

        note.tags.append(song_tag)
        note.tags.append(shift_tag)

        collection.add_note(note, DeckId(deck_id))
        card_ids = [card.id for card in note.cards()]
        if card_ids:
            collection.sched.suspend_cards(card_ids)


def parse_xsc_file(path: Path) -> XscDocument:
    if not path.exists():
        raise FileNotFoundError(f'File not found: {path}')

    content = path.read_text(encoding='utf-8')
    return parse_xsc_content(content)


def parse_xsc_content(content: str) -> XscDocument:
    sections = parse_sections(content)
    document = build_document(sections)
    return document


def parse_sections(content: str) -> List[Section]:
    sections: List[Section] = []
    lines = iter(content.splitlines())

    next(lines, None)
    next(lines, None)

    for line in lines:
        if not line:
            continue

        if line.startswith('SectionStart,'):
            name = line.removeprefix('SectionStart,')
            section_lines: List[Line] = []

            for section_line in lines:
                if section_line.startswith('SectionEnd,'):
                    break

                if not section_line:
                    continue

                if ',' in section_line:
                    key, value = section_line.split(',', 1)
                    section_lines.append(Line(key=key, value=value))

            sections.append(Section(name=name, lines=section_lines))

    return sections


def build_document(sections: Sequence[Section]) -> XscDocument:
    main_section = _find_section(sections, 'Main')
    loops_section = _find_section(sections, 'Loops')
    view_section = _find_optional_section(sections, 'View0')

    audio_file = parse_audio_file(main_section)
    loops = parse_loops(loops_section)
    tempo_adjustment = parse_tempo(view_section) if view_section else 1.0
    eq_settings = parse_eq(view_section) if view_section else None

    return XscDocument(
        audio_file=audio_file,
        loops=loops,
        tempo_adjustment=tempo_adjustment,
        eq_settings=eq_settings,
    )


def _find_section(sections: Sequence[Section], name: str) -> Section:
    for section in sections:
        if section.name == name:
            return section
    raise ValueError(f'Missing section: {name}')


def _find_optional_section(sections: Sequence[Section], name: str) -> Optional[Section]:
    for section in sections:
        if section.name == name:
            return section
    return None


def parse_audio_file(section: Section) -> AudioFile:
    sound_line = section.find_line('SoundFileName')
    info_line = section.find_line('SoundFileInfo')

    parts = sound_line.value.split(',')
    if len(parts) < 3:
        raise ValueError('Invalid SoundFileName line in .xsc')

    raw_path = parts[2]
    # `\C` is an escape character for `,` in the path
    raw_path = raw_path.replace(r'\C', ',')

    path = Path(raw_path)

    info_parts = info_line.value.split(',')
    if len(info_parts) < 5:
        raise ValueError('Invalid SoundFileInfo line in .xsc')

    channels = int(info_parts[2])
    sample_rate = int(info_parts[4])

    return AudioFile(path=path, sample_rate=sample_rate, channels=channels)


def parse_loops(section: Section) -> List[Loop]:
    loops: List[Loop] = []
    for line in section.lines:
        if line.key != 'L':
            continue
        loops.append(parse_loop_line(line.value))
    return loops


def parse_loop_line(value: str) -> Loop:
    parts = value.split(',')
    if len(parts) < 4:
        raise ValueError('Invalid loop line in .xsc')

    loop_id = int(parts[0])
    enabled = parts[1] == '1'
    start_samples = int(parts[2])
    duration_samples = int(parts[3])
    color = parts[5] if len(parts) > 5 else 'White'

    return Loop(
        id=loop_id,
        enabled=enabled,
        start_samples=start_samples,
        duration_samples=duration_samples,
        color=color,
    )


def parse_tempo(section: Section) -> float:
    line = section.find_line('FX_Speed')
    parts = line.value.split(',')
    if len(parts) < 3:
        raise ValueError('Invalid FX_Speed line in .xsc')
    tempo = float(parts[2])
    return tempo / 100000.0


def parse_eq(section: Section) -> Optional[List[float]]:
    try:
        line = section.find_line('FX_EQ')
    except ValueError:
        return None

    parts = line.value.split(',')
    if len(parts) < 4:
        return None

    eq_values = []
    for chunk in parts[3].split(':'):
        try:
            eq_values.append(float(chunk))
        except ValueError:
            continue
    return eq_values


def compute_hash(
    xsc_name: str,
    start_samples: int,
    end_samples: int,
    pitch_shift: float,
    speed: float,
) -> ChunkHash:
    sha = hashlib.sha256()
    sha.update(xsc_name.encode('utf-8'))
    sha.update(str(start_samples).encode('utf-8'))
    sha.update(str(end_samples).encode('utf-8'))
    sha.update(f'{pitch_shift:.6f}'.encode('utf-8'))
    sha.update(f'{speed:.6f}'.encode('utf-8'))
    return ChunkHash(sha.hexdigest())


def build_pitch_shift_sequence() -> List[PitchVariant]:
    sequence: List[PitchVariant] = [PitchVariant(shift=0.0, alterations=0)]
    steps_from_fifth: List[int] = [1, -1, 2, -2, 3, -3, 4, -4, 5, -5, 6]

    for steps in steps_from_fifth:
        raw = (steps * 7) % 12
        shift = ((raw + 6) % 12) - 6
        sequence.append(PitchVariant(shift=float(shift), alterations=abs(steps)))

    return sequence


def build_riff_name(xsc_path: Path, loop_id: int) -> str:
    seed = f'{xsc_path.stem}-{loop_id}'
    digest = hashlib.sha256(seed.encode('utf-8')).digest()

    color_word = _COLOR_WORDS[digest[0] % len(_COLOR_WORDS)]
    object_word = _OBJECT_WORDS[digest[1] % len(_OBJECT_WORDS)]
    animal_word = _ANIMAL_WORDS[digest[2] % len(_ANIMAL_WORDS)]

    return f'{color_word} {object_word} {animal_word}'


def text_color_for_shift(pitch_shift: float, alterations: int) -> Optional[str]:
    if alterations <= 0:
        return None

    palette_pos = ['#b65c5c', '#c24d4d', '#cd3f3f', '#d93030', '#e42222', '#ef1313']
    palette_neg = ['#4d6fb6', '#3f63c2', '#3157cd', '#234bd9', '#153ee4', '#0732ef']
    index = min(alterations, 6) - 1
    palette = palette_pos if pitch_shift > 0 else palette_neg
    return palette[index]


def wrap_with_color(text: str, color: Optional[str]) -> str:
    if not color or not text:
        return text
    return f'<span style="color:{color}">{text}</span>'


_COLOR_WORDS: tuple[str, ...] = (
    'Azure',
    'Brick',
    'Crimson',
    'Cyan',
    'Gold',
    'Indigo',
    'Ivory',
    'Lime',
    'Mauve',
    'Onyx',
    'Peach',
    'Plum',
    'Rose',
    'Sable',
    'Teal',
    'Umber',
)


_OBJECT_WORDS: tuple[str, ...] = (
    'Bell',
    'Brick',
    'Coin',
    'Drum',
    'Flute',
    'Glass',
    'Latch',
    'Loom',
    'Mask',
    'Pin',
    'Reel',
    'Rope',
    'Spoon',
    'Stone',
    'Torch',
    'Vial',
)


_ANIMAL_WORDS: tuple[str, ...] = (
    'Ant',
    'Bat',
    'Cat',
    'Crab',
    'Deer',
    'Dove',
    'Elk',
    'Fox',
    'Hare',
    'Jay',
    'Koi',
    'Lark',
    'Mole',
    'Newt',
    'Pika',
    'Stoat',
)


def resolve_mpv_path() -> Path:
    if sys.platform == 'darwin':
        mpvs = list(Path(sys.argv[0]).parent.parent.glob('**/mpv*'))
    elif sys.platform.startswith('win'):
        mpvs = list(Path(sys.argv[0]).parent.glob('**/mpv*.exe'))
    else:
        mpvs = list(Path(sys.argv[0]).parent.glob('**/mpv*'))

    for candidate in mpvs:
        if candidate.exists():
            print(f'Found mpv at {candidate}')
            return candidate

    raise RuntimeError('Could not locate the mpv binary shipped with Anki')


def export_chunk_with_mpv(
    mpv_path: Path,
    source_audio: Path,
    start_sec: float,
    end_sec: float,
    output_path: Path,
) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)

    cmd = [
        str(mpv_path),
        '--no-config',
        '--really-quiet',
        '--vid=no',
        f'--start={start_sec:.6f}',
        f'--end={end_sec:.6f}',
        '--of=ogg',
        '--oac=libopus',
        f'--o={output_path}',
        str(source_audio),
    ]

    try:
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as exc:  # noqa: BLE001
        raise RuntimeError(f'mpv failed to export audio: {exc}') from exc

    if not output_path.exists():
        raise RuntimeError('mpv did not produce an output file')


def build_song_tag(audio_path: Path) -> str:
    stem = audio_path.stem
    if not stem:
        return 'jam-track'

    collapsed = re.sub(r'[^A-Za-z0-9]+', ' ', stem).strip()
    if not collapsed:
        return 'jam-track'

    pieces = collapsed.split()
    slug = '-'.join(part.lower() for part in pieces)
    return f'jam-track-{slug}' if slug else 'jam-track'


def build_shift_tag(pitch_shift: float) -> str:
    epsilon = 0.001
    if pitch_shift > epsilon:
        direction = 'up'
    elif pitch_shift < -epsilon:
        direction = 'down'
    else:
        return 'jam-pitch-normal'

    magnitude = f'{abs(pitch_shift):.3f}'.rstrip('0').rstrip('.')
    if not magnitude:
        magnitude = '0'

    return f'jam-pitch-{direction}{magnitude}'


def build_media_filename(chunk_hash: str) -> str:
    return f'{chunk_hash}.ogg'


def split_artist_title(stem: str) -> Tuple[str, str]:
    if ' - ' in stem:
        artist, title = stem.split(' - ', 1)
        return artist.strip(), title.strip()
    return '', stem


def format_timestamp(seconds: float) -> str:
    total_ms = int(round(seconds * 1000))
    hours, remainder = divmod(total_ms // 1000, 3600)
    minutes, secs = divmod(remainder, 60)
    millis = total_ms % 1000
    return f'{hours:02}:{minutes:02}:{secs:02}.{millis:03}'