from __future__ import annotations

from dataclasses import replace
from bisect import bisect_left, bisect_right
from fractions import Fraction
import math

from generator import harmony
from generator.engraving import tablature
from generator.model import validate_document
from generator.engraving.attributes import barline, staff_attributes, tempo
from generator.engraving.fonts import Fonts
from generator.engraving.notation import ScoreIndex
from generator.engraving.notes import draw_event
from generator.engraving.spans import draw_beams, draw_lyric_spans, draw_spans, place_text_lanes
from generator.engraving.types import Anchor, EngravingError, Fragment, LayoutDocument, Page, Primitive, Style, System, bounds, line

def lyric_text(verse: dict) -> str:
  return ''.join(s['text'] for s in verse['syllables'] if s['kind'] in ('text', 'elision'))

def column(index: ScoreIndex, time: Fraction, start: Fraction, end: Fraction, fonts: Fonts, style: Style) -> tuple[list[Primitive], list[Anchor], list[Fragment]]:
  items, anchors, fragments = [], [], []
  events = index.events_at(time, include_sustained=time == start)
  grace_width = 0
  occupied = []
  for event in events:
    try:
      drawn, placed, pieces = draw_event(index, event, start, end, fonts, style)
    except ValueError as error:
      raise EngravingError(f'measure {event.measure["number"]}, voice {event.voice}, {event.value["id"]}: {error}') from error
    x = grace_width
    collidable = [p for p in drawn if p.role in ('head', 'tab', 'rest', 'slash', 'accidental', 'dot')]
    while any(p.box.moved(x, 0).overlaps(q.box, 2) for p in collidable for q in occupied):
      x += style.space * 1.4
    drawn = [p.moved(x, 0) for p in drawn]
    placed = [a.moved(x, 0) for a in placed]
    if event.value.get('grace'):
      grace_width = max((p.box.right for p in drawn), default=grace_width) + 9
    occupied.extend(p for p in drawn if p.role in ('head', 'tab', 'rest', 'slash', 'accidental', 'dot'))
    items.extend(drawn)
    items.extend(tablature.column(index, event, placed, fonts, style))
    anchors.extend(placed)
    fragments.extend(pieces)
    if event.start == time and not event.value.get('hidden'):
      verses = event.value.get('lyrics', [])
      for verse in verses:
        lane = index.lyric_lanes.index((event.voice, verse['number']))
        text = lyric_text(verse)
        if text:
          items.append(fonts.text(text, x + style.space / 2, lane * 23, 18, 'lyric', (event.value['id'] + ':lyric:' + verse['number'], event.value['id']), True))
  measure_start, _, measure = index.measure_at(time)
  for owner, _, item in index.harmonies_by_onset.get(time, []):
    if owner is measure:
      items.append(fonts.text(harmony.symbol(item), 0, 0, 21, 'harmony', (item['id'],)))
  for owner, i, value in index.tempos_by_onset.get(time, []):
    if owner is measure:
      items.extend(tempo(value, fonts, 0, 0, style.space, measure['id'] + f':tempo:{i}'))
  if time == measure_start:
    items.append(fonts.text(str(measure['number']), 0, -style.space * (index.state(time)['staff_details']['lines'] + 0.5), 12, 'measure-number', (measure['id'],)))
    if label := measure.get('section_label'):
      items.append(fonts.text(label, 0, 0, 19, 'section', (measure['id'] + ':section',)))
  return items, anchors, fragments

def make_system(index: ScoreIndex, start: Fraction, end: Fraction, fonts: Fonts, style: Style, justify: bool = False) -> System:
  items, anchors, fragments = [], [], []
  state = index.state(start)
  measure = index.measure_at(start)[2]
  header = staff_attributes(index, state, start, fonts, style.space, measure['id'] + ':initial', ('clef', 'key', 'meter', 'staff_details'))
  items.extend(header)
  cursor = max((p.box.right for p in header), default=0) + 22
  content_left = cursor - 9
  times = [start] + index.column_times[bisect_right(index.column_times, start):bisect_left(index.column_times, end)]
  placements = []
  groups = []
  for ti, time in enumerate(times):
    item_start, anchor_start = len(items), len(anchors)
    drawn, placed, pieces = column(index, time, start, end, fonts, style)
    leading = []
    measure_start, _, m = index.measure_at(time)
    if time == measure_start and time > start:
      previous = index.measure_at(time, before=True)[2]
      leading.extend(barline(previous.get('barlines', {}).get('right', 'regular'), 0, -style.space * (index.state(time)['staff_details']['lines'] - 1), previous['id']))
      old_state = index.state(time, before=True)
      fields = tuple(k for k in ('clef', 'key', 'meter', 'staff_details') if old_state[k] != index.state(time)[k])
      leading.extend(p.moved(12, 0) for p in staff_attributes(index, index.state(time), time, fonts, style.space, m['id'] + ':initial', fields))
    if time > start:
      for owner, ci, change in index.attributes_by_onset.get(time, []):
        extra = staff_attributes(index, index.state(time), time, fonts, style.space, owner['id'] + f':attributes:{ci}', tuple(change))
        leading.extend(p.moved(max((p.box.right for p in leading), default=0) + 10, 0) for p in extra)
    if leading:
      width = bounds(leading).right - bounds(leading).left + 14
      items.extend(p.moved(cursor - bounds(leading).left, 0) for p in leading)
      cursor += width
    box = bounds(drawn)
    x = cursor - min(0, box.left)
    placements.append((time, x))
    items.extend(p.moved(x, 0) for p in drawn)
    anchors.extend(a.moved(x, 0) for a in placed)
    fragments.extend(pieces)
    next_time = times[ti + 1] if ti + 1 < len(times) else end
    spring = style.column_gap + 13 * math.sqrt(min(4, float(next_time - time)))
    cursor = x + max(style.space * 1.4, box.right) + spring
    groups.append((item_start, len(items), anchor_start, len(anchors)))
  width = cursor + 6
  available = style.width - style.margin_x * 2
  if justify and width < available:
    stretch = (available - width) / max(1, len(times))
    for i, (first_item, last_item, first_anchor, last_anchor) in enumerate(groups):
      items[first_item:last_item] = [p.moved(i * stretch, 0) for p in items[first_item:last_item]]
      anchors[first_anchor:last_anchor] = [a.moved(i * stretch, 0) for a in anchors[first_anchor:last_anchor]]
    placements = [(time, x + i * stretch) for i, (time, x) in enumerate(placements)]
    width = available
  right = width - 5
  last_measure = index.measure_at(end, before=True)[2]
  tab_items = [p for p in items if p.role.startswith('tab')]
  items = [p for p in items if not p.role.startswith('tab')]
  items = draw_beams(index, anchors, items, start, end, content_left, right, style)
  items = draw_spans(index, anchors, items, start, end, content_left, right, fonts, style)
  items = place_text_lanes(items, fonts, style)
  items = draw_lyric_spans(index, items, anchors, start, end, content_left, right, fonts)
  changes = [(0, state['staff_details']['lines'])]
  for time, x in placements:
    lines = index.state(time)['staff_details']['lines']
    if changes[-1][1] != lines:
      changes.append((x - 8, lines))
  staff_lines = []
  for i, (x1, lines) in enumerate(changes):
    x2 = changes[i + 1][0] if i + 1 < len(changes) else width - 1
    for number in range(lines):
      y = -number * style.space
      staff_lines.append(line('staff', (), ((x1, y), (x2, y)), style.space * 0.13))
  items = staff_lines + items
  if end in index.boundaries:
    items.extend(barline(last_measure.get('barlines', {}).get('right', 'regular'), width - 1, -style.space * (index.state(end, before=True)['staff_details']['lines'] - 1), last_measure['id']))
  if start in index.boundaries:
    items.extend(barline(measure.get('barlines', {}).get('left', 'none'), 0, -style.space * (state['staff_details']['lines'] - 1), measure['id']))
  else:
    items.append(fonts.text(str(measure['number']) + ' (cont.)', 0, min(p.box.top for p in items) - 6, 12, 'continuation', (measure['id'],)))
  items = tablature.staff(index, items + tab_items, start, end, content_left, width - 1, fonts, style)
  return System(start, end, tuple(items), tuple(anchors), tuple(fragments), bounds(items))

def paginate(systems: list[System], document: dict, fonts: Fonts, style: Style) -> tuple[Page, ...]:
  pages = []
  placed = []
  header = []
  cursor = style.margin_y
  first = True

  def page_header(number):
    title = document['metadata']['title']
    size = 26 if number == 1 else 17
    item = fonts.text(title, style.margin_x, style.margin_y + size, size, 'title', ())
    while item.box.right > style.width - style.margin_x and size > 12:
      size -= 1
      item = fonts.text(title, style.margin_x, style.margin_y + size, size, 'title', ())
    if item.box.right > style.width - style.margin_x:
      raise EngravingError('score title cannot fit the page without clipping')
    result = [item]
    if number == 1:
      subtitle = ' - '.join(str(value) for value in (document['metadata'].get('composer'), document['staves'][0].get('name'), document['metadata'].get('style')) if value)
      if subtitle:
        result.append(fonts.text(subtitle, style.margin_x, item.box.bottom + 22, 16, 'subtitle', ()))
    return result

  for system in systems:
    height = system.box.bottom - system.box.top
    if first:
      header = page_header(len(pages) + 1)
      cursor = bounds(header).bottom + 24
      first = False
    if cursor + height > style.height - style.margin_y - 18:
      if not placed:
        raise EngravingError(f'at musical offset {system.start}: simultaneous notation needs {height:.1f}px vertically; fixed-size page cannot fit it')
      number = len(pages) + 1
      footer = fonts.text(str(number), style.width / 2, style.height - style.margin_y, 12, 'page-number', (), True)
      pages.append(Page(number, style.width, style.height, style.supersampling, tuple(placed), tuple(header + [p for s in placed for p in s.primitives] + [footer])))
      placed = []
      header = page_header(len(pages) + 1)
      cursor = bounds(header).bottom + 24
      if cursor + height > style.height - style.margin_y - 18:
        raise EngravingError(f'at musical offset {system.start}: simultaneous notation is too tall for a page')
    dx, dy = style.margin_x + 1, cursor - system.box.top
    moved = replace(system, primitives=tuple(p.moved(dx, dy) for p in system.primitives), anchors=tuple(a.moved(dx, dy) for a in system.anchors), box=system.box.moved(dx, dy))
    placed.append(moved)
    cursor += height + style.system_gap
  if placed:
    number = len(pages) + 1
    footer = fonts.text(str(number), style.width / 2, style.height - style.margin_y, 12, 'page-number', (), True)
    pages.append(Page(number, style.width, style.height, style.supersampling, tuple(placed), tuple(header + [p for s in placed for p in s.primitives] + [footer])))
  return tuple(pages)

def layout_score(document: dict, style: Style, fonts: Fonts) -> LayoutDocument:
  validate_document(document)
  index = ScoreIndex(document)
  if not index.measures:
    raise EngravingError('score has no measures')
  systems = []
  start = Fraction()
  available = style.width - style.margin_x * 2
  while start < index.end:
    candidates = index.boundary_times[bisect_right(index.boundary_times, start):]
    selected = None
    # Reserve the same title/credit and footer bands when considering a split,
    # so a wide pitch range can break vertically as well as horizontally.
    height_limit = style.height - style.margin_y * 2 - 110
    for end in candidates:
      trial = make_system(index, start, end, fonts, style)
      if trial.box.right - trial.box.left > available or trial.box.bottom - trial.box.top > height_limit:
        break
      selected = trial
    if selected is None:
      measure_end = candidates[0]
      for end in index.column_times[bisect_right(index.column_times, start):bisect_left(index.column_times, measure_end)]:
        trial = make_system(index, start, end, fonts, style)
        if trial.box.right - trial.box.left <= available and trial.box.bottom - trial.box.top <= height_limit:
          selected = trial
        else:
          break
    if selected is None:
      raise EngravingError(f'at musical offset {start}: irreducible simultaneous content exceeds the page width or height')
    systems.append(make_system(index, start, selected.end, fonts, style, justify=selected.end < index.end))
    start = selected.end
  layout = LayoutDocument(paginate(systems, document, fonts, style), style, fonts.fingerprint(style))
  from generator.engraving.audit import audit_layout
  audit_layout(document, layout, fonts)
  return layout