from __future__ import annotations

import array
import fractions
import hashlib
import json
import math
import pathlib
import sys
import wave

import generator.model
import generator.pitch

SAMPLE_RATE = 48000
PROFILE = 'soft_struck'
MAX_EVENTS = 1000000

class PlaybackUnavailable(ValueError):
  pass

def sounding_pitch(tone: dict, state: dict) -> fractions.Fraction:
  if 'unpitched' in tone:
    raise PlaybackUnavailable('unpitched notation has no generated_tone interpretation')
  value = generator.pitch.sounding_pitch(tone, state)
  frequency = 440 * 2 ** (float(value - 69) / 12)
  if not 20 <= frequency <= 18000:
    raise PlaybackUnavailable(f'sounding pitch {value} is outside the generated_tone range (20-18000 Hz)')
  return value

def compile_timeline(document: dict) -> dict:
  events, notes, linked = [], {}, {}
  absolute = fractions.Fraction()
  for measure in document['staves'][0]['measures']:
    for voice in measure['voices']:
      state = {key: measure[key] for key in ('transpose', 'clef', 'staff_details')}
      changes = iter(sorted(measure.get('attributes', []), key=lambda a: generator.model.fraction(a['onset'])))
      change = next(changes, None)
      for note in voice['notes']:
        onset = generator.model.fraction(note['onset'])
        while change is not None and generator.model.fraction(change['onset']) <= onset:
          state.update({key: change[key] for key in state if key in change})
          change = next(changes, None)
        notes[note['id']] = (note, absolute + onset, dict(state), voice['id'])
    absolute += generator.model.fraction(measure['duration'])
  adjustments = {}
  insertions = {}
  lanes = {}
  for item in notes.values():
    lanes.setdefault(item[3], []).append(item)
  for lane in lanes.values():
    pending = []
    previous = None
    for note, onset, state, voice in lane:
      if note.get('grace'):
        pending.append((note, onset, state, voice))
        continue
      if pending:
        cursor = onset
        specified = [n['grace'] for n, _, _, _ in pending]
        modes = {key for g in specified for key in ('make-time', 'steal-time-previous', 'steal-time-following') if key in g}
        if len(modes) > 1:
          raise PlaybackUnavailable('grace group mixes conflicting timing modes')
        mode = next(iter(modes), 'steal-time-following')
        if mode == 'make-time':
          lengths = [generator.model.fraction(g.get(mode, generator.model.rational(0))) for g in specified]
          total = sum(lengths, fractions.Fraction())
          insertions[onset] = max(insertions.get(onset, fractions.Fraction()), total)
        else:
          donor = previous[0] if mode == 'steal-time-previous' and previous else note if mode == 'steal-time-following' else None
          if donor is None:
            raise PlaybackUnavailable('grace notes have no preceding duration to steal')
          percentages = [generator.model.fraction(g[mode]) if mode in g else fractions.Fraction(25, len(pending)) for g in specified]
          if not 0 < sum(percentages) < 100:
            raise PlaybackUnavailable('grace group must steal between zero and 100 percent')
          lengths = [generator.model.fraction(donor['duration']) * amount / 100 for amount in percentages]
          total = sum(lengths, fractions.Fraction())
          delay, extension = adjustments.get(donor['id'], (fractions.Fraction(), fractions.Fraction()))
          if mode == 'steal-time-previous':
            cursor -= total
          else:
            delay += total
          extension -= total
          if generator.model.fraction(donor['duration']) + extension <= 0 or any(generator.model.fraction(p['duration']) + extension <= 0 for p in donor['pitches']):
            raise PlaybackUnavailable(f'{donor["id"]}: grace timing consumes a donor or chord tone')
          adjustments[donor['id']] = (delay, extension)
        for item, length in zip(pending, lengths):
          value, _, pitch_state, lane_id = item
          adjustments[value['id']] = (cursor - item[1], length)
          cursor += length
        pending = []
      previous = (note, onset)
    if pending:
      raise PlaybackUnavailable('trailing grace group has no following anchor')
  for identifier, (note, onset, state, voice) in notes.items():
    if note.get('cue'):
      continue
    delay, extension = adjustments.get(identifier, (fractions.Fraction(), fractions.Fraction()))
    performed = onset + delay + sum((length for time, length in insertions.items() if time < onset or time == onset and not note.get('grace')), fractions.Fraction())
    tremolo = note.get('tremolo')
    if tremolo and tremolo['type'] == 'stop':
      continue
    for pitch in note['pitches']:
      if pitch.get('cue'):
        continue
      length = generator.model.fraction(pitch['duration']) + extension
      if not note.get('grace'):
        length += sum((amount for time, amount in insertions.items() if onset < time < onset + generator.model.fraction(pitch['duration'])), fractions.Fraction())
      if length <= 0:
        raise PlaybackUnavailable(f'{identifier}: grace timing consumes a chord tone')
      value = sounding_pitch(pitch, state)
      if pitch.get('tie_from'):
        previous = linked.get(pitch['tie_from'])
        if previous is None or generator.model.fraction(previous['onset']) + generator.model.fraction(previous['duration']) != performed or generator.model.fraction(previous['pitch']) != value:
          raise PlaybackUnavailable(f'{identifier}: tie has inconsistent performed timing or pitch')
        previous['duration'] = generator.model.rational(generator.model.fraction(previous['duration']) + length)
        linked[pitch['id']] = previous
        continue
      if tremolo and tremolo['marks']:
        if pitch.get('tie_to'):
          raise PlaybackUnavailable('tied measured tremolo is not yet qualified for playback')
        written = pitch['written']
        step = min(fractions.Fraction(1), generator.model.TYPE_LENGTHS[written['type']]) * generator.model.fraction(written.get('ratio', generator.model.rational(1))) / 2 ** tremolo['marks']
        if tremolo['type'] == 'start':
          partner = notes[tremolo['partner']]
          if partner[1] != onset + generator.model.fraction(note['duration']) or len(partner[0]['pitches']) != len(note['pitches']):
            raise PlaybackUnavailable('paired tremolo requires contiguous equal-size chords')
          step *= 2
          length += generator.model.fraction(partner[0]['pitches'][note['pitches'].index(pitch)]['duration'])
        cursor, attack = fractions.Fraction(), 0
        while cursor < length:
          amount = value
          if tremolo['type'] == 'start' and attack % 2:
            other = partner[0]['pitches'][note['pitches'].index(pitch)]
            amount = sounding_pitch(other, partner[2])
          events.append({'id': f'{pitch["id"]}_{attack}', 'pitch': generator.model.rational(amount), 'onset': generator.model.rational(performed + cursor), 'duration': generator.model.rational(min(step, length - cursor))})
          cursor += step
          attack += 1
          if len(events) > MAX_EVENTS:
            raise PlaybackUnavailable('playback exceeds one million tone events')
      else:
        event = {'id': pitch['id'], 'pitch': generator.model.rational(value), 'onset': generator.model.rational(performed), 'duration': generator.model.rational(length)}
        events.append(event)
        linked[pitch['id']] = event
  if len(events) > MAX_EVENTS:
    raise PlaybackUnavailable('playback exceeds one million tone events')
  events.sort(key=lambda e: (generator.model.fraction(e['onset']), e['id']))
  return {'available': True, 'duration': generator.model.rational(absolute + sum(insertions.values(), fractions.Fraction())), 'events': events}

def write_pcm(path: pathlib.Path, values) -> str:
  samples = array.array('h', (round(max(-1, min(1, sample)) * 32767) for sample in values))
  if sys.byteorder != 'little':
    samples.byteswap()
  path.parent.mkdir(parents=True, exist_ok=True)
  with wave.open(str(path), 'wb') as output:
    output.setparams((1, 2, SAMPLE_RATE, 0, 'NONE', 'not compressed'))
    output.writeframes(samples.tobytes())
  return hashlib.sha256(path.read_bytes()).hexdigest()

def generate_tone(root: pathlib.Path, pitch: fractions.Fraction) -> dict:
  frequency = 440 * 2 ** (float(pitch - 69) / 12)
  period = fractions.Fraction(frequency / SAMPLE_RATE).limit_denominator(SAMPLE_RATE)
  loop_length, cycles = period.denominator, period.numerator
  attack = 12000
  name = hashlib.sha256(f'{PROFILE}:{pitch}'.encode('ascii')).hexdigest()[:24]
  relative = f'audio/{name}.wav'
  harmonics = [h for h in range(1, 7) if frequency * h < SAMPLE_RATE * 0.45]

  def sample(frame):
    phase = 2 * math.pi * cycles * frame / loop_length
    envelope = min(1, frame / 240) * (0.18 + 0.82 * (1 - min(1, frame / attack)) ** 3)
    return 0.3 * envelope * sum(math.sin(phase * h) / (h * h) for h in harmonics)

  checksum = write_pcm(root / relative, (sample(i) for i in range(attack + loop_length)))
  return {'id': name, 'path': relative, 'sha256': checksum, 'sample_rate': SAMPLE_RATE, 'frames': attack + loop_length, 'loop_start': attack, 'loop_end': attack + loop_length}

def compile_tree(root: pathlib.Path) -> None:
  tones = {}
  release = {'path': 'audio/release.wav', 'sample_rate': SAMPLE_RATE, 'frames': 2400}
  release['sha256'] = write_pcm(root / release['path'], ((1 - i / 2399) ** 2 for i in range(2400)))
  for path in sorted(root.rglob('*.json')):
    document = json.loads(path.read_text(encoding='utf-8'))
    try:
      playback = compile_timeline(document)
      used = {}
      for event in playback['events']:
        event['end'] = generator.model.rational(generator.model.fraction(event['onset']) + generator.model.fraction(event.pop('duration')))
        pitch = generator.model.fraction(event.pop('pitch'))
        if pitch not in tones:
          tones[pitch] = generate_tone(root, pitch)
        tone = tones[pitch]
        event['sample'] = tone['id']
        used[tone['id']] = tone
      playback.update({'samples': list(used.values()), 'release': release, 'profile': PROFILE})
    except PlaybackUnavailable as error:
      playback = {'available': False, 'reason': str(error)}
    document['playback'] = playback
    path.write_text(json.dumps(document, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')

def validate_assets(playback: dict, root: pathlib.Path, cache: dict | None = None) -> None:
  if not isinstance(playback, dict) or type(playback.get('available')) is not bool:
    raise ValueError('playback manifest requires boolean availability')
  if not playback['available']:
    if not isinstance(playback.get('reason'), str) or not playback['reason']:
      raise ValueError('unavailable playback must explain why')
    return
  cache = {} if cache is None else cache

  def check_sample(value):
    name = value['path']
    relative = pathlib.Path(name)
    if not name.startswith('audio/') or relative.is_absolute() or relative.suffix != '.wav' or any(p in ('', '.', '..') for p in name.split('/')) or '\\' in name or '\x00' in name:
      raise ValueError('invalid generated audio path')
    path = (root / relative).resolve()
    if not path.is_relative_to(root.resolve()):
      raise ValueError('generated audio escapes asset root')
    key = (str(path), value['sha256'], value['frames'], value['sample_rate'])
    if key in cache:
      return cache[key]
    if value['sample_rate'] != SAMPLE_RATE or not 0 < value['frames'] <= 96000:
      raise ValueError('invalid generated audio rate/length')
    if path.stat().st_size != 44 + value['frames'] * 2 or hashlib.sha256(path.read_bytes()).hexdigest() != value['sha256']:
      raise ValueError('generated audio checksum mismatch')
    with wave.open(str(path)) as audio:
      pcm = audio.readframes(value['frames'])
      if (audio.getnchannels(), audio.getsampwidth(), audio.getframerate(), audio.getnframes()) != (1, 2, SAMPLE_RATE, value['frames']) or len(pcm) != value['frames'] * 2:
        raise ValueError('invalid or truncated generated PCM WAV')
    samples = array.array('h')
    samples.frombytes(pcm)
    if sys.byteorder != 'little':
      samples.byteswap()
    if min(samples) < -32767:
      raise ValueError('generated PCM exceeds unit range')
    cache[key] = samples
    return samples

  duration = generator.model.fraction(playback['duration'])
  if duration <= 0 or len(playback['events']) > MAX_EVENTS:
    raise ValueError('invalid playback timeline length')
  release = check_sample(playback['release'])
  if not 1 < playback['release']['frames'] <= SAMPLE_RATE // 2:
    raise ValueError('invalid release envelope length')
  if release[0] != 32767 or release[-1] != 0 or any(a < b for a, b in zip(release, release[1:])):
    raise ValueError('release envelope must decrease from one to zero')
  if len(playback['samples']) > 512 or sum(s['frames'] for s in playback['samples']) > 16 * 1024 * 1024:
    raise ValueError('generated playback exceeds sample memory budget')
  samples = {}
  for sample in playback['samples']:
    if not sample['id'] or sample['id'] in samples or not 0 <= sample['loop_start'] < sample['loop_end'] == sample['frames']:
      raise ValueError('invalid sample identity or loop')
    check_sample(sample)
    samples[sample['id']] = sample
  previous = fractions.Fraction()
  identities = set()
  for event in playback['events']:
    onset, end = generator.model.fraction(event['onset']), generator.model.fraction(event['end'])
    if event['id'] in identities or event['sample'] not in samples or not previous <= onset < end <= duration:
      raise ValueError('invalid playback event identity, reference or timing')
    identities.add(event['id'])
    previous = onset