from __future__ import annotations

from fractions import Fraction
from math import lcm
from pathlib import Path
from xml.etree import ElementTree

from generator import harmony
from generator.model import alteration, fraction, validate_document

def element(parent, tag, content=None, **attributes):
  node = ElementTree.SubElement(parent, tag, {k.replace('_', '-'): str(v) for k, v in attributes.items() if v is not None})
  if content is not None:
    node.text = str(content)
  return node

def decimal(value) -> str:
  value = alteration(value) if isinstance(value, dict) else Fraction(value)
  denominator = value.denominator
  for prime in (2, 5):
    while denominator % prime == 0:
      denominator //= prime
  if denominator != 1:
    raise ValueError(f'MusicXML requires a terminating decimal here, not {value}')
  digits, remainder = [], abs(value.numerator) % value.denominator
  while remainder:
    digit, remainder = divmod(remainder * 10, value.denominator)
    digits.append(str(digit))
  return ('-' if value < 0 else '') + str(abs(value.numerator) // value.denominator) + ('.' + ''.join(digits) if digits else '')

def write_attributes(parent, state: dict, divisions=None) -> None:
  node = element(parent, 'attributes')
  if divisions is not None:
    element(node, 'divisions', divisions)
  if key := state.get('key'):
    item = element(node, 'key')
    if 'cancel' in key:
      element(item, 'cancel', key['cancel'])
    if key.get('custom'):
      steps = key.get('steps', [{'letter': letter, 'alter': amount} for letter, amount in zip('cdefgab', key['alterations']) if alteration(amount)])
      for step in steps:
        element(item, 'key-step', step['letter'].upper())
        element(item, 'key-alter', decimal(step['alter']))
        if 'accidental' in step:
          element(item, 'key-accidental', step['accidental'])
      for index, step in enumerate(steps):
        if 'octave' in step:
          element(item, 'key-octave', step['octave'], number=index + 1)
    else:
      element(item, 'fifths', key['fifths'])
      element(item, 'mode', key['mode'])
  if meter := state.get('meter'):
    item = element(node, 'time', symbol=meter.get('symbol', 'normal'))
    if meter.get('unmetered'):
      element(item, 'senza-misura')
    else:
      for group in meter['groups']:
        element(item, 'beats', '+'.join(map(str, group['beats'])))
        element(item, 'beat-type', group['unit'])
  if clef := state.get('clef'):
    item = element(node, 'clef')
    element(item, 'sign', clef['sign'])
    element(item, 'line', clef['line'])
    if clef['octave_change']:
      element(item, 'clef-octave-change', clef['octave_change'])
  if details := state.get('staff_details'):
    item = element(node, 'staff-details')
    element(item, 'staff-lines', details['lines'])
    for tuning in details.get('tunings', []):
      t = element(item, 'staff-tuning', line=tuning['line'])
      element(t, 'tuning-step', tuning['letter'].upper())
      element(t, 'tuning-alter', decimal(tuning['alter']))
      element(t, 'tuning-octave', tuning['octave'])
    if details.get('capo'):
      element(item, 'capo', details['capo'])
  transpose = state.get('transpose')
  if transpose and (transpose.get('explicit') or any(transpose[field] for field in ('diatonic', 'chromatic', 'octave_change'))):
    item = element(node, 'transpose')
    for field in ('diatonic', 'chromatic', 'octave_change'):
      element(item, field.replace('_', '-'), transpose[field])

def write_harmony(parent, value: dict, divisions: int) -> None:
  node = element(parent, 'harmony')
  if root := value.get('root'):
    item = element(node, 'root')
    element(item, 'root-step', root['letter'].upper())
    element(item, 'root-alter', decimal(root['alter']))
  element(node, 'kind', value['kind'], text=value.get('kind_text'))
  if 'inversion' in value:
    element(node, 'inversion', value['inversion'])
  if 'bass' in value:
    item = element(node, 'bass')
    element(item, 'bass-step', value['bass']['letter'].upper())
    element(item, 'bass-alter', decimal(value['bass']['alter']))
  for degree in value['degrees']:
    item = element(node, 'degree')
    amount = alteration(degree['alter'])
    if degree['type'] == 'alter':
      amount -= harmony.KINDS[value['kind']].get(degree['value'], 0)
    elif degree['type'] == 'add' and (degree['value'] - 1) % 7 == 6:
      amount += 1
    element(item, 'degree-value', degree['value'])
    element(item, 'degree-alter', decimal(amount))
    element(item, 'degree-type', degree['type'])
  element(node, 'offset', int(fraction(value['onset']) * divisions))

def write_note(parent, event: dict, voice: str, divisions: int, tuplets: list[dict]) -> None:
  if event['kind'] == 'spacer':
    item = element(parent, 'forward')
    element(item, 'duration', int(fraction(event['duration']) * divisions))
    element(item, 'voice', voice)
    return
  pitches = event['pitches'] or [None]
  for index, pitch in enumerate(pitches):
    printing = {'print-object': 'no'} if event.get('hidden') else {}
    if pitch:
      printing.update(pitch.get('print', {}))
    node = element(parent, 'note', **printing)
    if grace := event.get('grace'):
      element(node, 'grace', **{k: ('yes' if v else 'no') if isinstance(v, bool) else decimal(fraction(v) * divisions if k == 'make-time' else v) for k, v in grace.items()})
    if (pitch or event).get('cue', event.get('cue')):
      element(node, 'cue')
    if index:
      element(node, 'chord')
    if event['kind'] == 'rest':
      rest = element(node, 'rest', measure='yes' if event.get('measure_rest') else None)
      if position := event.get('rest_position'):
        element(rest, 'display-step', position['letter'].upper())
        element(rest, 'display-octave', position['octave'])
    elif pitch and 'unpitched' in pitch:
      unpitched = element(node, 'unpitched')
      element(unpitched, 'display-step', pitch['unpitched']['letter'].upper())
      element(unpitched, 'display-octave', pitch['unpitched']['octave'])
    else:
      item = element(node, 'pitch')
      element(item, 'step', pitch['letter'].upper() if pitch else 'B')
      element(item, 'alter', decimal(pitch['alter']) if pitch else '0')
      element(item, 'octave', pitch['octave'] if pitch else 4)
    if not event.get('grace'):
      element(node, 'duration', int(fraction(pitch['duration'] if pitch else event['duration']) * divisions))
    for marker in (pitch or {}).get('sound_ties', []):
      element(node, 'tie', type=marker['type'], time_only=','.join(map(str, marker['passes'])) if marker.get('passes') else None)
    for field, action in (('tie_from', 'stop'), ('tie_to', 'start')):
      if pitch and pitch.get(field):
        element(node, 'tie', type=action)
    if pitch and pitch.get('instrument'):
      element(node, 'instrument', id=pitch['instrument'])
    element(node, 'voice', voice)
    written = pitch['written'] if pitch else event['written']
    element(node, 'type', written['type'])
    for _ in range(written.get('dots', 0)):
      element(node, 'dot')
    if pitch and (display := pitch.get('accidental_display')):
      element(node, 'accidental', display['name'], **{k: v for k, v in display.items() if k != 'name'})
    if 'ratio' in written:
      ratio = fraction(written['ratio'])
      item = element(node, 'time-modification')
      element(item, 'actual-notes', ratio.denominator)
      element(item, 'normal-notes', ratio.numerator)
      if normal := written.get('normal'):
        element(item, 'normal-type', normal['type'])
        for _ in range(normal.get('dots', 0)):
          element(item, 'normal-dot')
    if stem := (pitch or event).get('stem', event.get('stem')):
      element(node, 'stem', stem)
    head = pitch.get('notehead') if pitch else None
    if head:
      element(node, 'notehead', head['name'], **{k: v for k, v in head.items() if k != 'name'})
    elif event.get('notehead'):
      element(node, 'notehead', event['notehead'])
    if index == 0:
      for beam in event.get('beams', []):
        element(node, 'beam', beam['type'], number=beam['level'])
    notations = element(node, 'notations')
    for action in (pitch or {}).get('printed_ties', []):
      element(notations, 'tied', type=action)
    for field, action in (('tie_from', 'stop'), ('tie_to', 'start')):
      if pitch and pitch.get(field):
        element(notations, 'tied', type=action)
    if index == 0:
      for span in tuplets:
        for endpoint, action in (('start', 'start'), ('end', 'stop')):
          if span[endpoint] != event['id']:
            continue
          item = element(notations, 'tuplet', type=action, number=span['depth'] + 1, bracket=None if span['bracket'] == 'auto' else span['bracket'], show_number=span['show_number'], show_type=span.get('show_type', 'none'), line_shape=span.get('line_shape', 'straight'))
          if action == 'start':
            for tag, field in (('tuplet-actual', 'note_count'), ('tuplet-normal', 'space_count')):
              detail = span.get(tag.replace('-', '_'), {})
              group = element(item, tag)
              element(group, 'tuplet-number', span[field])
              if detail.get('type'):
                element(group, 'tuplet-type', detail['type'])
              for _ in range(detail.get('dots', 0)):
                element(group, 'tuplet-dot')
      if tremolo := event.get('tremolo'):
        element(element(notations, 'ornaments'), 'tremolo', tremolo['marks'], type=tremolo['type'])
    if pitch and (tab := pitch.get('tablature')):
      technical = element(notations, 'technical')
      element(technical, 'string', tab['string'])
      element(technical, 'fret', tab['fret'])
    if not len(notations):
      node.remove(notations)
    if index == 0:
      for verse in event.get('lyrics', []):
        lyric = element(node, 'lyric', number=verse['number'])
        for syllable in verse['syllables']:
          element(lyric, syllable['kind'], syllable['text'], **syllable['attributes'])

def write_barline(bar, measure, location):
  if not any(location in measure.get(field, {}) for field in ('barlines', 'repeats', 'endings')):
    return
  node = element(bar, 'barline', location=location)
  element(node, 'bar-style', measure.get('barlines', {}).get(location, 'regular'))
  if ending := measure.get('endings', {}).get(location):
    element(node, 'ending', ending.get('text', ''), number=','.join(map(str, ending['numbers'])), type=ending['type'], print_object='yes' if ending.get('visible', True) else 'no')
  if repeat := measure.get('repeats', {}).get(location):
    element(node, 'repeat', direction=repeat['direction'], times=repeat.get('times'))

def musicxml_tree(document: dict) -> ElementTree.ElementTree:
  validate_document(document)
  root = ElementTree.Element('score-partwise', {'version': '4.0'})
  metadata, staff = document['metadata'], document['staves'][0]
  element(element(root, 'work'), 'work-title', metadata['title'])
  identification = element(root, 'identification')
  element(identification, 'creator', metadata.get('composer', ''), type='composer')
  miscellaneous = element(identification, 'miscellaneous')
  element(miscellaneous, 'miscellaneous-field', metadata.get('style', 'Medium Swing'), name='vamp:style')
  if metadata.get('transpose_origin'):
    from generator.model import pitch_name
    element(miscellaneous, 'miscellaneous-field', pitch_name(metadata['transpose_origin']), name='vamp:tonic')
  defaults = element(root, 'defaults')
  element(defaults, 'music-font', font_family='Petaluma')
  element(defaults, 'word-font', font_family='Petaluma Script')
  element(defaults, 'lyric-font', font_family='Petaluma Script')
  definition = element(element(root, 'part-list'), 'score-part', id='P1')
  element(definition, 'part-name', staff['name'])
  for instrument in staff.get('instruments', []):
    item = element(definition, 'score-instrument', id=instrument['id'])
    element(item, 'instrument-name', instrument['name'])
    if instrument['abbreviation']:
      element(item, 'instrument-abbreviation', instrument['abbreviation'])
    if instrument['sound']:
      element(item, 'instrument-sound', instrument['sound'])
    if instrument.get('solo'):
      element(item, 'solo')
    if 'ensemble' in instrument:
      element(item, 'ensemble', instrument['ensemble'])
    if 'virtual' in instrument:
      virtual = element(item, 'virtual-instrument')
      for tag, value in instrument['virtual'].items():
        element(virtual, tag, value)
  for instrument in staff.get('instruments', []):
    if instrument.get('midi'):
      item = element(definition, 'midi-instrument', id=instrument['id'])
      for tag, value in instrument['midi'].items():
        element(item, tag, value)
  part = element(root, 'part', id='P1')
  absolute = Fraction()
  previous_state = {}
  for index, measure in enumerate(staff['measures']):
    bar = element(part, 'measure', number=index + 1, implicit='yes' if measure['implicit'] else 'no')
    write_barline(bar, measure, 'left')
    times = [fraction(measure['duration'])]
    times.extend(fraction(n[field]) for v in measure['voices'] for n in v['notes'] for field in ('onset', 'duration'))
    times.extend(fraction(p['duration']) for v in measure['voices'] for n in v['notes'] for p in n['pitches'])
    times.extend(fraction(item['onset']) for field in ('attributes', 'harmonies', 'tempos') for item in measure.get(field, []))
    times.extend(fraction(s[field]) - absolute for s in staff.get('octave_shifts', []) for field in ('start', 'end'))
    divisions = lcm(*(t.denominator for t in times))
    state = {k: measure[k] for k in ('key', 'meter', 'clef', 'transpose', 'staff_details') if k in measure}
    write_attributes(bar, {k: v for k, v in state.items() if previous_state.get(k) != v}, divisions)
    previous_state.update(state)
    cursor = Fraction()
    for change in measure.get('attributes', []):
      onset = fraction(change['onset'])
      if onset != cursor:
        element(element(bar, 'forward' if onset > cursor else 'backup'), 'duration', int(abs(onset - cursor) * divisions))
      write_attributes(bar, change)
      previous_state.update({k: v for k, v in change.items() if k != 'onset'})
      cursor = onset
    if cursor:
      element(element(bar, 'backup'), 'duration', int(cursor * divisions))
    for value in measure['harmonies']:
      write_harmony(bar, value, divisions)
      if value.get('chord_optional'):
        direction = element(bar, 'direction', print_object='no')
        element(element(direction, 'direction-type'), 'words', 'vamp:optional-chord')
        element(direction, 'offset', int(fraction(value['onset']) * divisions))
    if measure.get('section_label'):
      element(element(element(bar, 'direction'), 'direction-type'), 'rehearsal', measure['section_label'], halign='left')
    for tempo in measure.get('tempos', []):
      direction = element(bar, 'direction')
      metronome = element(element(direction, 'direction-type'), 'metronome')
      element(metronome, 'beat-unit', tempo['written']['type'])
      for _ in range(tempo['written'].get('dots', 0)):
        element(metronome, 'beat-unit-dot')
      element(metronome, 'per-minute', decimal(tempo['bpm']))
      element(direction, 'offset', int(fraction(tempo['onset']) * divisions))
    for si, span in enumerate(staff.get('octave_shifts', [])):
      for field in ('start', 'end'):
        onset = fraction(span[field]) - absolute
        inside = 0 <= onset < fraction(measure['duration']) if field == 'start' else 0 < onset <= fraction(measure['duration'])
        if not inside:
          continue
        direction = element(bar, 'direction')
        element(element(direction, 'direction-type'), 'octave-shift', type='stop' if field == 'end' else 'up' if span['octaves'] > 0 else 'down', size=abs(span['octaves']) * 7 + 1, number=si + 1)
        element(direction, 'offset', int(onset * divisions))
    cursor = Fraction()
    for voice in measure['voices']:
      if cursor:
        element(element(bar, 'backup'), 'duration', int(cursor * divisions))
      cursor = Fraction()
      for event in voice['notes']:
        onset = fraction(event['onset'])
        if onset > cursor:
          element(element(bar, 'forward'), 'duration', int((onset - cursor) * divisions))
        write_note(bar, event, voice['id'], divisions, staff.get('tuplets', []))
        cursor = onset + fraction(event['duration'])
      if cursor < fraction(measure['duration']):
        element(element(bar, 'forward'), 'duration', int((fraction(measure['duration']) - cursor) * divisions))
        cursor = fraction(measure['duration'])
    write_barline(bar, measure, 'right')
    absolute += fraction(measure['duration'])
  return ElementTree.ElementTree(root)

def write_musicxml(path: Path, document: dict) -> None:
  tree = musicxml_tree(document)
  ElementTree.indent(tree, space='  ')
  tree.write(path, encoding='utf-8', xml_declaration=True)