from __future__ import annotations

import dataclasses
import fractions
import typing

import generator.arrange.ending_plan
import generator.arrange.pickup
import generator.model
import generator.score

@dataclasses.dataclass(frozen=True)
class ArrangementSection:
  label: str
  source_measure_index: int
  measure_count: int
  section_index: int
  role: typing.Literal["accompaniment", "melody", "vocal", "rest", "slash"]

def phrase_aware_edge_measure_count(body_measure_count: int) -> int:
  if body_measure_count >= 8:
    return 8
  if body_measure_count >= 4:
    return 4
  return body_measure_count

def section_role(part_type: str, section_index: int) -> typing.Literal["accompaniment", "melody", "vocal", "rest", "slash"]:
  if part_type == "acoustic":
    return "slash" if section_index == 3 else "accompaniment"
  if part_type == "electric":
    if section_index == 4:
      return "slash"
    if section_index == 2:
      return "melody"
    if section_index in {1, 3, 5}:
      return "accompaniment"
    return "rest"
  return "vocal" if section_index in {1, 5} else "rest"

def full_arrangement_sections(measures: list[dict], part_type: str) -> list[ArrangementSection]:
  if not measures:
    return []
  body_start = generator.score.body_start_measure_index(measures)
  if body_start >= len(measures):
    return []
  body_measure_count = len(measures) - body_start
  edge_measure_count = phrase_aware_edge_measure_count(body_measure_count)
  edge_source_measure = body_start + body_measure_count - edge_measure_count
  specs = (
    ("Rhythm Lead-in", edge_source_measure, edge_measure_count, 0),
    ("Vocal Melody", body_start, body_measure_count, 1),
    ("Guitar Melody", body_start, body_measure_count, 2),
    ("Acoustic Solo", body_start, body_measure_count, 3),
    ("Electric Solo", body_start, body_measure_count, 4),
    ("Vocal Melody", body_start, body_measure_count, 5),
  )
  return [
    ArrangementSection(label=label, source_measure_index=start, measure_count=count, section_index=index, role=section_role(part_type, index))
    for label, start, count, index in specs
    if count > 0
  ]

@dataclasses.dataclass(frozen=True)
class FormPlan:
  sections: tuple[ArrangementSection, ...]
  windows: tuple[tuple[fractions.Fraction, fractions.Fraction, str], ...]
  stops: tuple[tuple[int, fractions.Fraction] | None, ...]
  @classmethod
  def read(cls, document, part):
    source = generator.score.source_staff(document)['measures']
    sections = tuple(full_arrangement_sections(source, part))
    stops = tuple(section_stop(source, section, sections[i + 1] if i + 1 < len(sections) else None) for i, section in enumerate(sections))
    return cls(sections, _windows(document, part, sections), stops)

def section_stop(source, section, following):
  if following is None or not generator.score.body_start_measure_index(source):
    return None
  last = section.source_measure_index + section.measure_count - 1
  pad = generator.arrange.pickup.offset(source[last], source[0])
  if following.role == 'slash':
    return last, pad
  if following.role == 'melody' and generator.arrange.pickup.enters(following, source):
    attacks = [generator.model.fraction(note['onset']) for note in generator.score.measure_notes(source[0]) if note['pitches']]
    if attacks:
      return last, pad + min(attacks)
  return None

def windows(document, part):
  return FormPlan.read(document, part).windows

def _windows(document, part, sections):
  source = generator.score.source_staff(document)['measures']
  result, cursor = [], fractions.Fraction()
  for section in sections:
    duration = sum((generator.model.fraction(m['duration']) for m in source[section.source_measure_index:section.source_measure_index + section.measure_count]), fractions.Fraction())
    role = {'slash': 'solo', 'vocal': 'melody'}.get(section.role, section.role)
    if section.role in ('vocal', 'melody') and generator.arrange.pickup.enters(section, source) and result:
      attacks = [generator.model.fraction(n['onset']) for n in generator.score.measure_notes(source[0]) if n['pitches']]
      if attacks:
        start = cursor - generator.model.fraction(source[0]['duration']) + min(attacks)
        a, _, old = result[-1]
        result[-1] = (a, start, old)
        result.append((start, cursor, part + '_' + role))
    end = cursor + duration
    start = cursor
    if section.role == 'slash' and generator.score.body_start_measure_index(source) and result:
      # Allocate the source's complete pickup bar before composition. This
      # does not inspect either generated solo; a role may leave it silent.
      start = cursor - generator.model.fraction(source[0]['duration'])
      a, _, old = result[-1]
      result[-1] = (a, start, old)
    result.append((start, end, part + '_' + role))
    cursor = end
  result.append((cursor, cursor + generator.arrange.ending_plan.duration(source), part + ('_accompaniment' if part != 'vocal' else '_rest')))
  return tuple(result)