from __future__ import annotations

import copy
import fractions
import pathlib

import generator.accompaniment
import generator.arrange.accompaniment
import generator.acoustic
import generator.acoustic_voicing
import generator.arrange.ending
import generator.arrange.form
import generator.arrange.pickup
import generator.electric
import generator.fretboard
import generator.guitar
import generator.melody
import generator.model
import generator.musicxml
import generator.parts
import generator.score
import generator.solo.arrangement
import generator.comping
import generator.vocabulary
import generator.production

FINAL_SECTION_INDEX = 5

def arrange_document(document: dict, part_type: str) -> dict:
  """Canonical source-only template; solos are inserted by production roles."""
  document = generator.parts.prepare_document(document)
  with generator.vocabulary.using(generator.vocabulary.load(), part_type + '_accompaniment'):
    backing = generator.comping.Canonical(document, part_type, 'canonical') if part_type != 'vocal' else None
    return _arrange_document(document, part_type, accompaniment=backing)

def section_source(source: list[dict], section, next_section=None) -> list[dict]:
  result = copy.deepcopy(source)
  for measure in result:
    measure['harmonies'] = [h for h in measure['harmonies'] if section.section_index != FINAL_SECTION_INDEX or not h.get('chord_optional')]
    for chord in measure['harmonies']:
      chord.pop('chord_optional', None)
  if generator.arrange.pickup.enters(next_section, source):
    destination = result[section.source_measure_index + section.measure_count - 1]
    generator.arrange.pickup.harmonies(destination, result[0])
  return result

def _arrange_document(document: dict, part_type: str, plans=None, *, owned=None, accompaniment=None) -> dict:
  if part_type not in generator.parts.PART_TYPES:
    raise ValueError(f'unknown arrangement part: {part_type}')
  staff = generator.score.source_staff(document)
  source = staff['measures']
  sections = generator.arrange.form.full_arrangement_sections(source, part_type)
  output, tuplets, octave_shifts, voicing_origins = [], [], [], []
  absolute = fractions.Fraction()
  source_starts = []
  source_time = fractions.Fraction()
  for m in source:
    source_starts.append(source_time)
    source_time += generator.model.fraction(m['duration'])
  for index, section in enumerate(sections):
    role = section.role if owned is None or section.role in owned else 'rest'
    section_measures = section_source(source, section, sections[index + 1] if index + 1 < len(sections) else None)
    previous = generator.accompaniment.carried_chords(section_measures)
    segment = []
    stop = None
    next_section = sections[index + 1] if index + 1 < len(sections) else None
    if next_section and generator.score.body_start_measure_index(source):
      last = section.source_measure_index + section.measure_count - 1
      if next_section.role == 'slash':
        stop = (last, generator.model.fraction(section_measures[last]['duration']) - generator.model.fraction(source[0]['duration']))
      elif next_section.role == 'melody' and generator.arrange.pickup.enters(next_section, source):
        attacks = [generator.model.fraction(n['onset']) for n in generator.score.measure_notes(source[0]) if n['pitches']]
        if attacks:
          stop = (last, generator.model.fraction(section_measures[last]['duration']) - generator.model.fraction(source[0]['duration']) + min(attacks))
    for si in range(section.source_measure_index, section.source_measure_index + section.measure_count):
      m = section_measures[si]
      if role == 'accompaniment' and accompaniment is not None:
        result = accompaniment(m, previous[si], si, sources=section_measures, bounds=(section.source_measure_index, section.source_measure_index + section.measure_count), stop=stop)
        result['id'] = f'measure_{len(output) + len(segment) + 1}'
        result['section_label'] = section.label if not segment else ''
      else:
        result = generator.parts.build_measure(m, f'measure_{len(output) + len(segment) + 1}', role, part_type, previous[si], section.label if not segment else '')
      segment.append(result)
      if part_type == 'electric' and section.role in ('melody', 'accompaniment'):
        voicing_origins.append(generator.electric.MeasureOrigin(source_starts[si], section.role, result))
    segment_length = sum((generator.model.fraction(m['duration']) for m in segment), fractions.Fraction())
    if role in ('vocal', 'melody'):
      start = source_starts[section.source_measure_index]
      for shift in staff.get('octave_shifts', []):
        low, high = max(start, generator.model.fraction(shift['start'])), min(start + segment_length, generator.model.fraction(shift['end']))
        if low < high:
          octave_shifts.append({**shift, 'id': f'section_{section.section_index}_{shift["id"]}', 'start': generator.model.rational(absolute + low - start), 'end': generator.model.rational(absolute + high - start)})
    pickup = None
    if role in ('vocal', 'melody') and generator.arrange.pickup.enters(section, source) and output:
      pickup = generator.parts.build_measure(source[0], f'pickup_{section.section_index}', section.role, part_type, None)
    tuplets.extend(generator.score.reidentify(([pickup] if pickup else []) + segment, staff.get('tuplets', [])))
    if pickup:
      destination = output[-1]
      pad = generator.arrange.pickup.offset(destination, pickup)
      onset = generator.arrange.pickup.splice(destination, pickup)
      if part_type == 'electric':
        voicing_origins = [generator.electric.MeasureOrigin(o.onset, o.role, o.measure, o.start, onset if o.measure is destination else o.end) for o in voicing_origins]
        voicing_origins.append(generator.electric.MeasureOrigin(-pad, 'melody', destination, onset))
      for shift in staff.get('octave_shifts', []):
        low, high = max(onset - pad, generator.model.fraction(shift['start'])), min(generator.model.fraction(pickup['duration']), generator.model.fraction(shift['end']))
        if low < high:
          offset = absolute - generator.model.fraction(pickup['duration'])
          octave_shifts.append({**shift, 'id': f'pickup_{section.section_index}_{shift["id"]}', 'start': generator.model.rational(offset + low), 'end': generator.model.rational(offset + high)})
    output.extend(segment)
    absolute += segment_length
  if sections and sections[-1].section_index == FINAL_SECTION_INDEX:
    generator.arrange.ending.append(source, output, part_type,
      sounding=part_type != 'vocal' and (owned is None or 'accompaniment' in owned),
      center=generator.vocabulary.role(part_type + '_accompaniment').get('center', 55))
  tuplets = generator.arrange.pickup.retain_links(output, tuplets)
  result = generator.score.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 == 'acoustic':
    if accompaniment is None and (owned is None or 'accompaniment' in owned):
      generator.acoustic.prepare(result)
      generator.acoustic_voicing.place(result)
    generator.acoustic.validate_comping(result)
    generator.acoustic.combine(result)
  elif part_type == 'electric' and (owned is None or 'melody' in owned):
    generator.melody.harmonize(result)
    generator.guitar.place_guitar(result)
    if generator.accompaniment.connect_continuations(result):
      generator.guitar.place_guitar(result)
  # Source labels identify the head while arranging; published labels identify
  # the complete performance, including its lead-in and repeated choruses.
  for number, measure in enumerate(output, 1):
    measure['number'] = str(number)
  generator.fretboard.validate_document(result)
  return result

def generate_performance_scores(source_dir: pathlib.Path, output_dir: pathlib.Path, *, composition_jobs=2, cache_dir=None, progress=False, timings=None) -> int:
  return generator.production.generate_tree(source_dir, output_dir, jobs=composition_jobs, cache_dir=cache_dir, progress=progress, timings=timings)