from __future__ import annotations

import copy
import fractions
import pathlib

import generator.accompaniment
import generator.acoustic
import generator.acoustic_voicing
import generator.arrange.form
import generator.arrange.pickup
import generator.electric
import generator.guitar
import generator.melody
import generator.model
import generator.musicxml
import generator.parts
import generator.rhythm
import generator.scales.library
import generator.scales.mapping
import generator.scales.timeline
import generator.score
import generator.solo.arrangement
import generator.solo.sequences

FINAL_SECTION_INDEX = 6

def arrange_document(document: dict, part_type: str, *, rhythms: generator.rhythm.Library, scales: generator.scales.library.Library | None = None) -> dict:
  """Place the regular form; the song pipeline inserts solos after all heads exist."""
  document = generator.parts.prepare_document(document)
  plans = plan_arrangement(document, rhythms) if part_type != 'vocal' else {}
  return _arrange_document(document, part_type, plans, scales=scales)

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 plan_arrangement(document: dict, rhythms: generator.rhythm.Library) -> dict[int, generator.accompaniment.Plan]:
  source = generator.score.source_staff(document)['measures']
  result = {}
  sections = generator.arrange.form.full_arrangement_sections(source, 'acoustic')
  for index, section in enumerate(sections):
    if section.role == 'slash':
      continue
    measures = section_source(source, section, sections[index + 1] if index + 1 < len(sections) else None)
    previous = generator.accompaniment.carried_chords(measures)
    start = section.source_measure_index
    result[section.section_index] = generator.accompaniment.plan_section(measures[start:start + section.measure_count], rhythms, initial=previous[start], electric=section.section_index in (1, 5, 6))
  return result

def _arrange_document(document: dict, part_type: str, plans: dict[int, generator.accompaniment.Plan], *, scales: generator.scales.library.Library | None = None, reports: list[str] | None = None, source_name: str = '') -> 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 = [], [], [], []
  solos = []
  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):
    section_measures = section_source(source, section, sections[index + 1] if index + 1 < len(sections) else None)
    previous = generator.accompaniment.carried_chords(section_measures)
    segment = []
    for si in range(section.source_measure_index, section.source_measure_index + section.measure_count):
      m = section_measures[si]
      rhythm = None
      if section.role == 'accompaniment':
        plan = plans[section.section_index]
        rhythm = (plan.acoustic_bass[len(segment)], plan.acoustic_upper[len(segment)]) if part_type == 'acoustic' else plan.electric[len(segment)]
      result = generator.parts.build_measure(m, f'measure_{len(output) + len(segment) + 1}', section.role, part_type, previous[si], section.label if not segment else '', rhythm)
      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 section.role == 'slash':
      solos.append(generator.scales.timeline.Solo(absolute, absolute + segment_length, section.label))
    if section.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 section.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
  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':
    generator.acoustic.prepare(result)
    generator.acoustic_voicing.place(result)
    generator.acoustic.combine(result)
  elif part_type == 'electric':
    generator.melody.harmonize(result)
    voicing_links = generator.electric.link_voicings(source, voicing_origins)
    generator.guitar.place_guitar(result)
    if generator.accompaniment.connect_continuations(result):
      generator.guitar.place_guitar(result)
    generator.electric.reuse_voicings(voicing_links)
  if scales and solos:
    report = generator.scales.mapping.annotate(result, solos, scales, source_name or document['metadata']['title'])
    if reports is not None:
      reports.append(report)
  return result

def generate_performance_scores(source_dir: pathlib.Path, output_dir: pathlib.Path, *, rhythms: generator.rhythm.Library, scales: generator.scales.library.Library, sequences: generator.solo.sequences.Library) -> int:
  count = 0
  for path in generator.parts.source_paths(source_dir):
    try:
      document = generator.parts.prepare_document(generator.musicxml.load_musicxml(path))
      plans = plan_arrangement(document, rhythms)
    except ValueError as error:
      raise ValueError(f'{path}: {error}') from error
    parts = {}
    for part in generator.parts.PART_TYPES:
      try:
        parts[part] = _arrange_document(document, part, plans)
      except ValueError as error:
        raise ValueError(f'{path}, part {part}: {error}') from error
    try:
      report = generator.solo.arrangement.apply(document, parts, rhythms=rhythms, scales=scales, sequences=sequences)
    except ValueError as error:
      raise ValueError(f'{path}: {error}') from error
    for part, result in parts.items():
      generator.score.write_json(output_dir / f'{generator.score.slugify(path.stem)}_{part}.json', result)
      count += 1
    report_path = output_dir.parent / 'analysis' / f'{generator.score.slugify(path.stem)}.txt'
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.write_text(document['metadata']['title'] + '\nSource: ' + path.name + '\n\n' + report, encoding='utf-8')
  return count