#!/usr/bin/env python3
"""
Extended iMelody-like player.

Supports a richer subset of the grammar you provided:

- Octave prefix: "*0" .. "*8" placed immediately before a note (sets octave)
- Notes: basic notes a-g, sharps with leading '#' (e.g. #c), flats with leading '&' (e.g. &d)
- Duration digits: 0..5. Mapping (assumption):
        1 -> whole note (4 quarters)
        2 -> half note (2 quarters)
        3 -> quarter note (1 quarter)
        4 -> eighth note (1/2 quarter)
        5 -> sixteenth note (1/4 quarter)
        0 -> thirty-second note (1/8 quarter)
    (This maps larger digit -> shorter durations; choice is an inferred assumption.)
- Duration specifiers: '.' (dotted = *1.5), ':' (double = *2.0), ';' (staccato = *0.75)
    (These meanings are reasonable in musical contexts; if you want different semantics tell me.)
- Rests: 'r' followed by duration digit and optional specifier
- Volume: V0..V15 sets absolute volume; V+ / V- adjust up/down by 1
- Repeat: ( ... @N [V+/V-]) where N=0 means repeat forever (capped to 16 iterations here)
- LED / vibe / backlight tokens: 'ledon', 'ledoff', 'vibeon', 'vibeoff', 'backon', 'backoff'

Notes on assumptions:
- I cap infinite repeats (N=0) to 16 iterations to avoid infinite loops. We can change this.
- Octave mapping: default octave is 4 (C4 == MIDI 60). '*n' shifts octave accordingly.
- Volume maps linearly to amplitude: amplitude = 0.8 * (volume / 15).

Usage:
    python3 play_imelody.py "c4c4c4r2..." --tempo 120 --outfile out.wav
    python3 play_imelody.py "(c4r4@3)V8c4" --play
"""
import argparse
import math
import re
import struct
import wave
import subprocess
from pathlib import Path


NOTE_TO_MIDI = {
    'c': 60,
    'd': 62,
    'e': 64,
    'f': 65,
    'g': 67,
    'a': 69,
    'b': 71,
}


def midi_to_freq(midi):
    # A4 = MIDI 69 = 440 Hz
    return 440.0 * (2 ** ((midi - 69) / 12.0))


def _duration_quarters_from_digit(digit: int) -> float:
    # Assumption mapping described above.
    if digit == 0:
        return 4.0 / 32.0
    if digit >= 1:
        return 4.0 / (2 ** (digit - 1))
    raise ValueError("invalid duration digit")


class ParseError(Exception):
    pass


def parse_sequence(seq):
    """Parse the extended sequence into a flat list of events.

    Returns a list of dicts with type: 'note','rest','cmd' where note has
    midi, duration_quarters, volume_at_note, etc.
    """
    s = seq.strip()

    # tokenization helpers
    i = 0
    n = len(s)

    # parser state
    current_volume = 8  # default 0..15
    default_octave = 4

    def parse_items(limit=None):
        nonlocal i, current_volume, default_octave
        items = []
        while i < n:
            ch = s[i]
            if ch.isspace():
                i += 1
                continue
            if ch == '(':
                # find matching ')' via recursion: parse inside then expect @count
                i += 1
                start = i
                inner = parse_items()
                # when parse_items returns due to seeing ')' it will have advanced i to after ')'
                # now parse @count
                if i < n and s[i] == '@':
                    i += 1
                    m = re.match(r"(\d+)", s[i:])
                    if not m:
                        raise ParseError("Expected repeat count after @")
                    rep = int(m.group(1))
                    i += len(m.group(1))
                    # optional volume modifier
                    volmod = None
                    if i < n and s[i] in ('V'):
                        # V+ / V-
                        if s[i:i+2] in ('V+','V-'):
                            volmod = s[i:i+2]
                            i += 2
                    # expand repeats (cap infinite 0 -> 16)
                    cap = 16
                    if rep == 0:
                        rep = cap
                    for r in range(rep):
                        if volmod == 'V+':
                            current_volume = min(15, current_volume + 1)
                        elif volmod == 'V-':
                            current_volume = max(0, current_volume - 1)
                        # deep copy inner with current volume
                        for it in inner:
                            copy = dict(it)
                            # if note without explicit volume, set volume
                            if copy.get('type') in ('note','rest') and 'volume' not in copy:
                                copy['volume'] = current_volume
                            items.append(copy)
                    continue
                else:
                    raise ParseError("Missing @count after repeat group")
            if ch == ')':
                # end of current group
                i += 1
                return items

            # volume token Vn or V+ / V-
            if ch == 'V':
                # absolute or modifier
                if s[i:i+2] in ('V+','V-'):
                    if s[i+1] == '+':
                        current_volume = min(15, current_volume + 1)
                    else:
                        current_volume = max(0, current_volume - 1)
                    i += 2
                    continue
                m = re.match(r"V(\d{1,2})", s[i:])
                if m:
                    v = int(m.group(1))
                    current_volume = max(0, min(15, v))
                    i += 1 + len(m.group(1))
                    continue
                else:
                    raise ParseError("Invalid volume token")

            # LED/vibe/backlight commands
            for cmd in ('ledon','ledoff','vibeon','vibeoff','backon','backoff'):
                if s.startswith(cmd, i):
                    items.append({'type':'cmd','cmd':cmd})
                    i += len(cmd)
                    break
            else:
                # note/rest parsing
                octave = default_octave
                if s[i] == '*':
                    # octave prefix like *4
                    i += 1
                    if i < n and s[i].isdigit():
                        octave = int(s[i])
                        i += 1
                    else:
                        raise ParseError('Invalid octave prefix')
                # accidentals
                accidental = 0
                if i < n and s[i] in ('#','&'):
                    if s[i] == '#':
                        accidental = 1
                    else:
                        accidental = -1
                    i += 1

                if i >= n:
                    break
                ch2 = s[i].lower()
                if ch2 == 'r':
                    # rest
                    i += 1
                    if i < n and s[i].isdigit():
                        d = int(s[i])
                        i += 1
                    else:
                        raise ParseError('Expected duration digit after rest')
                    dur = _duration_quarters_from_digit(d)
                    # optional specifier
                    mult = 1.0
                    if i < n and s[i] in ('.',':',';'):
                        spec = s[i]
                        i += 1
                        if spec == '.':
                            mult = 1.5
                        elif spec == ':':
                            mult = 2.0
                        else:
                            mult = 0.75
                    dur *= mult
                    items.append({'type':'rest','duration_quarters':dur,'volume':current_volume})
                    continue
                if ch2 in NOTE_TO_MIDI:
                    note = ch2
                    i += 1
                    if i < n and s[i].isdigit():
                        d = int(s[i])
                        i += 1
                    else:
                        raise ParseError('Expected duration digit after note')
                    dur = _duration_quarters_from_digit(d)
                    mult = 1.0
                    if i < n and s[i] in ('.',':',';'):
                        spec = s[i]
                        i += 1
                        if spec == '.':
                            mult = 1.5
                        elif spec == ':':
                            mult = 2.0
                        else:
                            mult = 0.75
                    dur *= mult
                    # compute midi number
                    midi = NOTE_TO_MIDI[note] + (octave - 4) * 12 + accidental
                    items.append({'type':'note','midi':midi,'duration_quarters':dur,'volume':current_volume})
                    continue
                # unknown token
                raise ParseError(f"Unexpected token at {i}: '{s[i:]}'")
        return items

    parsed = parse_items()
    return parsed


def synthesize(items, tempo=120, sample_rate=44100):
    """Return PCM16 bytes for parsed sequence.

    Each item is a dict with type 'note'/'rest' and fields as created by parser.
    """
    quarter = 60.0 / tempo
    frames = []

    for it in items:
        if it['type'] == 'rest':
            dur_q = it['duration_quarters']
            duration = quarter * dur_q
            nframes = int(duration * sample_rate)
            frames.extend([0] * nframes)
            continue
        if it['type'] == 'cmd':
            # non-audio event; ignore in audio stream but keep for logs
            continue
        if it['type'] == 'note':
            dur_q = it['duration_quarters']
            duration = quarter * dur_q
            nframes = int(duration * sample_rate)
            midi = it['midi']
            freq = midi_to_freq(midi)
            vol = it.get('volume', 8)
            amplitude = 0.8 * (max(0, min(15, vol)) / 15.0)
            max_amp = int(32767 * amplitude)
            for i in range(nframes):
                t = i / sample_rate
                v = math.sin(2.0 * math.pi * freq * t)
                frames.append(int(v * max_amp))
            continue
        # ignore unknown
    pcm = struct.pack('<' + 'h' * len(frames), *frames) if frames else b''
    return pcm, sample_rate


def write_wav(path: Path, pcm_bytes: bytes, sample_rate: int):
    with wave.open(str(path), 'wb') as w:
        w.setnchannels(1)
        w.setsampwidth(2)  # 16-bit
        w.setframerate(sample_rate)
        w.writeframes(pcm_bytes)


def try_play(path: Path):
    # Try a few common players; prefer aplay on Linux
    for cmd in (['aplay', str(path)], ['paplay', str(path)]):
        try:
            subprocess.run(cmd, check=True)
            return True
        except FileNotFoundError:
            continue
        except subprocess.CalledProcessError:
            # player returned non-zero
            return False
    print("No suitable player found (tried aplay/paplay). WAV saved.")
    return False


def main():
    p = argparse.ArgumentParser()
    p.add_argument('sequence', help='compact iMelody-like sequence (e.g. c4c4r2...)')
    p.add_argument('--tempo', type=int, default=120, help='BPM tempo (quarter note)')
    p.add_argument('--outfile', default='out_imelody.wav', help='WAV output file')
    p.add_argument('--play', action='store_true', help='Attempt to play the WAV after generation')
    args = p.parse_args()

    items = parse_sequence(args.sequence)
    pcm, sr = synthesize(items, tempo=args.tempo)
    outp = Path(args.outfile)
    write_wav(outp, pcm, sr)
    print(f"Wrote WAV: {outp} ({len(pcm)} bytes, {sr} Hz)")
    if args.play:
        try_play(outp)


if __name__ == '__main__':
    main()