from __future__ import annotations

import copy
import warnings
from fractions import Fraction
from pathlib import Path

from generator.arrange.form import full_arrangement_sections
from generator.arrange.fretboard import pitches_for_voicing
from generator.model import beat_lengths, fraction, key_label, make_event, meter_label, rational
from generator.musicxml import load_musicxml
from generator.transposition import concert_document, shift_pitch
from generator.score import copy_source_voice, make_arranged_document, make_measure, measure_notes, reidentify, slugify, source_staff, write_json

PART_TYPES = ('acoustic','electric','vocal')
FINAL_SECTION_INDEX = 6


def propagated_chords(measures: list[dict], *, include_optional: bool = True) -> list[dict | None]:
  result = []
  current = None
  for measure in measures:
    result.append(current)
    for h in measure.get('harmonies', []):
      if include_optional or not h.get('chord_optional'):
        current = h
  return result


def active_chord(measure: dict, onset: Fraction, previous: dict | None) -> dict | None:
  active = previous
  for h in measure['harmonies']:
    if fraction(h['onset']) > onset:
      break
    active = h
  return active if active and active['kind'] != 'none' else None


def generated_voice(measure: dict, role: str, previous: dict | None) -> list[dict]:
  length = fraction(measure['duration'])
  if role == 'rest':
    return [make_event('rest', Fraction(), length, 'rest', written={'type':'whole','dots':0}) | {'measure_rest':True}]
  if role == 'rhythm' and measure['time_signature']=='4/4' and not measure.get('attributes'):
    attacks = [(Fraction(n,2), Fraction(1,2)) for n in (0,3,4,5,7)]
  elif role == 'accompaniment' and measure['time_signature']=='4/4' and not measure.get('attributes'):
    attacks = [(Fraction(1,2),Fraction(1,2))]
  else:
    attacks = []
    states = [(Fraction(), measure['meter'])]
    states.extend((fraction(a['onset']), a['meter']) for a in measure.get('attributes', []) if 'meter' in a)
    for index, (onset, meter) in enumerate(states):
      end = states[index+1][0] if index+1 < len(states) else length
      beats = beat_lengths(meter)
      beat_index = 0
      while onset < end:
        beat = min(beats[beat_index % len(beats)] if beats else end-onset, end-onset)
        duration = beat if role == 'slash' or not beats else min(beat/2, Fraction(1,2))
        attacks.append((onset, duration))
        onset += beat
        beat_index += 1
  events = []
  cursor = Fraction()
  for onset, duration in attacks:
    if onset >= length:
      break
    if onset > cursor:
      events.append(make_event(f'rest-{len(events)}',cursor,onset-cursor,'rest'))
    duration = min(duration,length-onset)
    chord = active_chord(measure,onset,previous)
    source_attack = any(fraction(n['onset'])==onset and n['kind']=='note' and not n.get('grace') and any(not p.get('tie_from') for p in n['pitches']) for n in measure_notes(measure))
    if role=='slash':
      event = make_event(f'slash-{len(events)}',onset,duration,'note')
      event['notehead'] = 'slash'
    elif not chord or role=='accompaniment' and source_attack:
      event = make_event(f'rest-{len(events)}',onset,duration,'rest')
    else:
      pitches = pitches_for_voicing(chord)
      event = make_event(f'chord-{len(events)}',onset,duration,'note',pitches)
      if not pitches:
        warnings.warn(f'{measure["id"]}: no playable voicing for {chord["symbol"]}; using labelled rhythmic slashes', stacklevel=2)
        event['notehead'] = 'slash'
    events.append(event)
    cursor = onset + duration
  if cursor < length:
    events.append(make_event(f'rest-{len(events)}',cursor,length-cursor,'rest'))
  return events


def arrange_document(document: dict, part_type: str) -> dict:
  if part_type not in PART_TYPES:
    raise ValueError(f'unknown arrangement part: {part_type}')
  document = concert_document(document)
  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)
      notes = copy_source_voice(m,strip_lyrics=section.role=='melody')['notes'] if section.role in ('vocal','melody') else generated_voice(m,section.role,previous[si])
      if section.role == 'melody':
        for note in notes:
          note['pitches'] = [shift_pitch(p,7,12) for p in note['pitches']]
      result = make_measure(f'measure-{len(output)+len(segment)+1}',m,notes,section_label=section.label if not segment else '')
      if part_type!='vocal':
        result['clef'] = {'sign':'G','line':2,'octave_change':0}
        result['transpose'] = {'diatonic':0,'chromatic':0,'octave_change':-1}
        for change in result.get('attributes',[]):
          change.pop('clef',None)
          change.pop('transpose',None)
      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 = copy.deepcopy(source[0])
      destination = output[-1]
      pad = fraction(destination['duration']) - fraction(pickup['duration'])
      if pad < 0:
        raise ValueError('pickup is longer than the preceding arrangement measure')
      notes = copy_source_voice(pickup,strip_lyrics=False)['notes']
      for n in notes:
        n['onset'] = rational(fraction(n['onset'])+pad)
      if pad:
        notes.insert(0,make_event('pickup-rest',Fraction(),pad,'rest'))
      destination['voices'][0]['notes'] = notes
      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)
  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('*.musicxml')):
    document = load_musicxml(path)
    for part in PART_TYPES:
      write_json(output_dir/f'{slugify(path.stem)}-{part}.json',arrange_document(document,part))
      count += 1
  return count