from __future__ import annotations

import copy
from fractions import Fraction
from pathlib import Path

from generator.arrange.form import full_arrangement_sections
from generator.guitar import place_guitar
from generator.parts import PART_TYPES, build_measure, propagated_chords, validate_source
from generator.model import fraction, key_label, make_event, meter_label, rational
from generator.musicxml import load_musicxml
from generator.pitch import concert_document
from generator.score import copy_source_voices, make_arranged_document, reidentify, slugify, source_staff, write_json

FINAL_SECTION_INDEX = 6

def arrange_document(document: dict, part_type: str) -> dict:
  if part_type not in PART_TYPES:
    raise ValueError(f'unknown arrangement part: {part_type}')
  validate_source(document)
  document = concert_document(document, octave_transposing=True)
  staff = source_staff(document)
  source = staff['measures']
  sections = full_arrangement_sections(source, part_type)
  output, tuplets, octave_shifts = [], [], []
  absolute = Fraction()
  source_starts = []
  source_time = Fraction()
  for m in source:
    source_starts.append(source_time)
    source_time += fraction(m['duration'])
  for section in sections:
    previous = propagated_chords(source, include_optional=section.section_index != FINAL_SECTION_INDEX)
    segment = []
    for si in range(section.source_measure_index, section.source_measure_index + section.measure_count):
      m = copy.deepcopy(source[si])
      m['harmonies'] = [h for h in m['harmonies'] if section.section_index != FINAL_SECTION_INDEX or not h.get('chord_optional')]
      for h in m['harmonies']:
        h.pop('chord_optional', None)
      result = build_measure(m, f'measure-{len(output) + len(segment) + 1}', section.role, part_type, previous[si], section.label if not segment else '')
      segment.append(result)
    identity_segment = segment
    segment_length = sum((fraction(m['duration']) for m in segment), Fraction())
    if section.role in ('vocal', 'melody'):
      start = source_starts[section.source_measure_index]
      for shift in staff.get('octave_shifts', []):
        low, high = max(start, fraction(shift['start'])), min(start + segment_length, fraction(shift['end']))
        if low < high:
          octave_shifts.append({**shift, 'id': f'section-{section.section_index}-{shift["id"]}', 'start': rational(absolute + low - start), 'end': rational(absolute + high - start)})
    if part_type == 'vocal' and section.section_index in (1, 5) and section.source_measure_index == 1 and source[0].get('implicit') and output:
      pickup = build_measure(source[0], source[0]['id'], 'vocal', part_type, None)
      destination = output[-1]
      pad = fraction(destination['duration']) - fraction(pickup['duration'])
      if pad < 0:
        raise ValueError('pickup is longer than the preceding arrangement measure')
      voices = copy_source_voices(pickup, strip_lyrics=False)
      for voice in voices:
        for n in voice['notes']:
          n['onset'] = rational(fraction(n['onset']) + pad)
        if pad:
          voice['notes'].insert(0, make_event(f'pickup-rest-{voice["id"]}', Fraction(), pad, 'rest') | {'voice': voice['id']})
      destination['voices'] = voices
      identity_segment = [destination, *segment]
      if pickup['harmonies']:
        destination['harmonies'] = [h for h in destination['harmonies'] if fraction(h['onset']) < pad]
        for h in pickup['harmonies']:
          h['onset'] = rational(fraction(h['onset']) + pad)
          destination['harmonies'].append(h)
      state = {k: copy.deepcopy(pickup[k]) for k in ('key', 'clef', 'meter', 'transpose')}
      if pad:
        destination['attributes'].append({'onset': rational(pad), **state})
      else:
        destination.update(state)
        destination['key_signature'] = key_label(state['key'])
        destination['time_signature'] = meter_label(state['meter'])
      for change in pickup.get('attributes', []):
        destination['attributes'].append({**change, 'onset': rational(fraction(change['onset']) + pad)})
      for shift in staff.get('octave_shifts', []):
        low, high = fraction(shift['start']), min(fraction(pickup['duration']), fraction(shift['end']))
        if low < high:
          offset = absolute - fraction(pickup['duration'])
          octave_shifts.append({**shift, 'id': f'pickup-{section.section_index}-{shift["id"]}', 'start': rational(offset + low), 'end': rational(offset + high)})
    tuplets.extend(reidentify(identity_segment, staff.get('tuplets', [])))
    output.extend(segment)
    absolute += segment_length
  result = make_arranged_document(document, staff, staff_id='staff-arranged-' + part_type, staff_name=part_type.capitalize(), role='vocal' if part_type == 'vocal' else 'guitar', arrangement_part_type=part_type, measures=output, tuplets=tuplets, octave_shifts=octave_shifts)
  result['metadata'].setdefault('tempo', 124)
  if part_type != 'vocal':
    place_guitar(result)
  return result

def generate_performance_scores(source_dir: Path, output_dir: Path) -> int:
  if not source_dir.exists():
    raise ValueError(f'performance source directory not found: {source_dir}')
  count = 0
  for path in sorted(source_dir.glob('*.xml')):
    document = load_musicxml(path)
    for part in PART_TYPES:
      try:
        write_json(output_dir / f'{slugify(path.stem)}-{part}.json', arrange_document(document, part))
      except ValueError as error:
        raise ValueError(f'{path}, part {part}: {error}') from error
      count += 1
  return count