from __future__ import annotations

import copy
from dataclasses import dataclass
from fractions import Fraction
import re

from generator.model import alteration, fraction, key_label, meter_label, rational

MAX_DEPTH = 16
MAX_MEASURES = 10000

def pass_numbers(value: str) -> list[int]:
  if not re.fullmatch(r'[1-9][0-9]*(?:,\s*[1-9][0-9]*)*', value.strip()):
    raise ValueError('ending/time-only numbers must be comma-separated positive integers')
  result = [int(item) for item in value.split(',')]
  if len(set(result)) != len(result) or max(result) > MAX_MEASURES:
    raise ValueError('duplicate or excessive ending/time-only number')
  return sorted(result)

@dataclass
class Ending:
  start: int
  end: int
  numbers: list[int]
  text: str
  closed: bool
  visible: bool

@dataclass
class Region:
  start: int
  end: int
  extent: int
  count: int
  endings: list[Ending]

@dataclass(frozen=True)
class Occurrence:
  measure: int
  passes: tuple[int, ...]

def analyze(measures: list[dict]) -> tuple[list[Region], list[Ending]]:
  regions, endings, starts = [], [], []
  active = None
  for index, measure in enumerate(measures):
    for side, boundary in (('left', index), ('right', index + 1)):
      if marker := measure.get('endings', {}).get(side):
        numbers = marker['numbers']
        if numbers != pass_numbers(','.join(map(str, numbers))):
          raise ValueError(f'measure {index + 1}: invalid ending passes')
        if marker['type'] == 'start':
          if active is not None:
            raise ValueError(f'measure {index + 1}: overlapping numbered endings')
          active = (boundary, marker)
        elif marker['type'] in ('stop', 'discontinue'):
          if active is None or active[1]['numbers'] != numbers or active[0] >= boundary:
            raise ValueError(f'measure {index + 1}: ending stop has no matching nonempty start')
          first, value = active
          endings.append(Ending(first, boundary, numbers, value.get('text') or ', '.join(map(str, numbers)) + '.', marker['type'] == 'stop', value.get('visible', True)))
          active = None
        else:
          raise ValueError(f'measure {index + 1}: invalid ending type')
      if marker := measure.get('repeats', {}).get(side):
        direction = marker['direction']
        if direction == 'forward' and side == 'left':
          starts.append(boundary)
          if len(starts) > MAX_DEPTH:
            raise ValueError('repeat nesting exceeds 16 levels')
        elif direction == 'backward' and side == 'right':
          first = starts.pop() if starts else 0
          count = marker.get('times', 2)
          if type(count) is not int or not 1 <= count <= MAX_MEASURES or first >= boundary:
            raise ValueError(f'measure {index + 1}: invalid repeat count or empty repeat')
          regions.append(Region(first, boundary, boundary, count, []))
        else:
          raise ValueError(f'measure {index + 1}: repeat direction disagrees with its boundary')
  if starts or active is not None:
    raise ValueError('unterminated repeat or numbered ending')
  assigned = set()
  for region in regions:
    first = next((ending for ending in endings if region.start <= ending.start < region.end and ending.end == region.end), None)
    if first is None:
      continue
    group = [first]
    while following := next((ending for ending in endings if ending.start == group[-1].end), None):
      group.append(following)
    numbers = [number for ending in group for number in ending.numbers]
    count = max(numbers)
    if sorted(numbers) != list(range(1, count + 1)) or count not in group[-1].numbers or len(group) < 2:
      raise ValueError(f'measure {region.start + 1}: endings must cover each pass exactly once and finish with the last pass')
    marker = measures[region.end - 1]['repeats']['right']
    if 'times' in marker and marker['times'] != count:
      raise ValueError(f'measure {region.end}: repeat count contradicts ending passes')
    region.count, region.extent, region.endings = count, group[-1].end, group
    assigned.update(id(ending) for ending in group)
  if any(id(ending) not in assigned for ending in endings):
    raise ValueError('numbered ending is not attached to a repeat group')
  for a, region in enumerate(regions):
    for other in regions[a + 1:]:
      if region.start == other.start and region.extent == other.extent:
        raise ValueError('duplicate repeat region')
      left, right = sorted((region, other), key=lambda item: (item.start, -item.extent))
      if left.start < right.start < left.extent < right.extent:
        raise ValueError('crossing repeat regions are ambiguous')
      for ending in left.endings:
        if right.start < ending.start < right.extent or right.start < ending.end < right.extent:
          raise ValueError('nested repeat crosses an ending boundary')
  return regions, endings

def performed_order(measures: list[dict]) -> list[Occurrence]:
  regions, _ = analyze(measures)
  result = []

  def walk(start, end, parents, passes, excluded=()):
    if len(parents) > MAX_DEPTH:
      raise ValueError('repeat nesting exceeds 16 levels')
    index = start
    while index < end:
      skipped = next((ending for ending in excluded if ending.start == index), None)
      if skipped:
        index = skipped.end
        continue
      candidates = [r for r in regions if r.start == index and r.extent <= end and id(r) not in parents]
      if candidates:
        region = max(candidates, key=lambda r: r.extent)
        for number in range(1, region.count + 1):
          walk(index, region.extent, (*parents, id(region)), (*passes, number), [e for e in region.endings if number not in e.numbers])
        index = region.extent
      else:
        result.append(Occurrence(index, passes))
        if len(result) > MAX_MEASURES:
          raise ValueError('expanded repeat traversal exceeds 10000 measures')
        index += 1

  walk(0, len(measures), (), ())
  return result

def expand_document(document: dict) -> dict:
  staff = document['staves'][0]
  source = staff['measures']
  if not any(m.get('repeats') or m.get('endings') for m in source):
    return copy.deepcopy(document)
  order = performed_order(source)
  output = copy.deepcopy(document)
  target = output['staves'][0]
  target['measures'], target['tuplets'], target['octave_shifts'] = [], [], []
  starts = [Fraction()]
  harmonies_before = []
  active_harmony = None
  event_measures = {n['id']: i for i, m in enumerate(source) for v in m['voices'] for n in v['notes']}
  event_ids = {(oi, n['id']): f'occurrence-{oi + 1}-n{ni + 1}' for oi, o in enumerate(order) for ni, n in enumerate(n for v in source[o.measure]['voices'] for n in v['notes'])}
  for measure in source:
    starts.append(starts[-1] + fraction(measure['duration']))
    harmonies_before.append(active_harmony)
    for harmony in measure['harmonies']:
      active_harmony = harmony
  absolute = Fraction()
  open_ties = {}
  expected_ties, used_ties = set(), set()
  previous_state = {}
  previous_index = -1
  active_tuplets = {}
  for oi, occurrence in enumerate(order):
    si = occurrence.measure
    measure = copy.deepcopy(source[si])
    measure['id'] = f'occurrence-{oi + 1}'
    measure.pop('repeats', None)
    measure.pop('endings', None)
    if active_harmony and oi and si <= previous_index and not any(fraction(h['onset']) == 0 for h in measure['harmonies']):
      restored = copy.deepcopy(harmonies_before[si]) if harmonies_before[si] else {'kind': 'none', 'symbol': 'N.C.', 'degrees': []}
      restored['onset'] = rational(0)
      measure['harmonies'].insert(0, restored)
    for side, marker in source[si].get('repeats', {}).items():
      measure.get('barlines', {}).pop(side, None)
    if oi and si > previous_index:
      for field, value in previous_state.items():
        if field not in measure.get('initial_attributes', {}):
          measure[field] = copy.deepcopy(value)
    measure.pop('initial_attributes', None)
    if oi and measure['key'].get('cancel') is not None:
      measure['key']['cancel'] = previous_state['key'].get('fifths', 0)
    measure['key_signature'] = key_label(measure['key'])
    measure['time_signature'] = meter_label(measure['meter'])
    previous_state = {field: copy.deepcopy(measure[field]) for field in ('key', 'meter', 'clef', 'transpose', 'staff_details')}
    for change in measure.get('attributes', []):
      previous_state.update({field: copy.deepcopy(value) for field, value in change.items() if field != 'onset'})
    event_map, tone_ids = {}, {}
    for hi, harmony in enumerate(measure['harmonies']):
      harmony['id'] = f'{measure["id"]}-h{hi + 1}'
    for ni, note in enumerate(n for v in measure['voices'] for n in v['notes']):
      old = note['id']
      note['id'] = f'{measure["id"]}-n{ni + 1}'
      event_map[old] = note['id']
      for pi, pitch in enumerate(note['pitches']):
        original = pitch['id']
        pitch['id'] = f'{note["id"]}-p{pi + 1}'
        tone_ids[pitch['id']] = original
        if 'sound_ties' not in pitch:
          pitch['sound_ties'] = [{'type': action} for field, action in (('tie_from', 'stop'), ('tie_to', 'start')) if pitch.get(field)]
        expected_ties.update((original, marker['type']) for marker in pitch['sound_ties'])
        pitch.pop('tie_from', None)
        pitch.pop('tie_to', None)
    for voice in measure['voices']:
      for note in voice['notes']:
        onset = absolute + fraction(note['onset'])
        for pitch in note['pitches']:
          markers = pitch.pop('sound_ties', [])
          pitch.pop('printed_ties', None)
          unpitched = pitch.get('unpitched', {})
          identity = (voice['id'], pitch.get('letter'), alteration(pitch.get('alter', 0)), pitch.get('octave'), unpitched.get('letter'), unpitched.get('octave'), pitch.get('instrument'))
          for marker in sorted(markers, key=lambda t: t['type'] == 'start'):
            if marker.get('passes') and (not occurrence.passes or occurrence.passes[-1] not in marker['passes']):
              continue
            used_ties.add((tone_ids[pitch['id']], marker['type']))
            if marker['type'] == 'stop':
              pending = open_ties.pop(identity, None)
              if not pending or pending[1] != onset:
                raise ValueError(f'measure {source[si]["number"]}, passes {occurrence.passes}: tie stop has no adjacent start on the performed path; use time-only for pass-specific ties')
              old, _ = pending
              old['tie_to'], pitch['tie_from'] = pitch['id'], old['id']
            else:
              if identity in open_ties:
                raise ValueError(f'measure {source[si]["number"]}: overlapping tie starts')
              open_ties[identity] = (pitch, onset + fraction(pitch['duration']))
        if partner := note.get('tremolo', {}).get('partner'):
          pi = oi + event_measures[partner] - si
          low, high = sorted((oi, pi))
          if pi < 0 or pi >= len(order) or any(order[j].measure != si + j - oi for j in range(low, high + 1)):
            raise ValueError(f'measure {source[si]["number"]}: repeat jump splits a paired tremolo')
          note['tremolo']['partner'] = event_ids[pi, partner]
    continued = {}
    for span in staff.get('tuplets', []):
      members = [n for v in measure['voices'] for n in v['notes'] if span['id'] in n.get('tuplet_ids', [])]
      if members:
        previous = active_tuplets.get(span['id']) if si == previous_index + 1 else None
        if previous and previous['continued_end'] and span['start'] not in event_map:
          rendered = previous
          rendered['end'] = members[-1]['id']
          rendered['continued_end'] = span['end'] not in event_map
        else:
          rendered = {**copy.deepcopy(span), 'id': f'{measure["id"]}-{span["id"]}', 'start': members[0]['id'], 'end': members[-1]['id'], 'continued_start': span['start'] not in event_map, 'continued_end': span['end'] not in event_map}
          target['tuplets'].append(rendered)
        identifier = rendered['id']
        continued[span['id']] = rendered
        for note in members:
          note['tuplet_ids'] = [identifier if t == span['id'] else t for t in note['tuplet_ids']]
    active_tuplets = continued
    for shift in staff.get('octave_shifts', []):
      low, high = max(starts[si], fraction(shift['start'])), min(starts[si + 1], fraction(shift['end']))
      if low < high:
        target['octave_shifts'].append({**shift, 'id': f'{measure["id"]}-{shift["id"]}', 'start': rational(absolute + low - starts[si]), 'end': rational(absolute + high - starts[si])})
    target['measures'].append(measure)
    absolute += fraction(measure['duration'])
    previous_index = si
  if open_ties:
    raise ValueError('tie start has no matching endpoint on the performed path; use time-only for pass-specific ties')
  if expected_ties != used_ties:
    raise ValueError('time-only tie refers to a pass that is never performed')
  return output