from __future__ import annotations

import copy
import json
import re
from pathlib import Path

from generator.model import SCHEMA_VERSION, validate_document


def load_json(path: Path) -> dict:
  document = json.loads(path.read_text(encoding='utf-8'))
  validate_document(document)
  return document


def write_json(path: Path, document: dict) -> None:
  validate_document(document)
  path.parent.mkdir(parents=True, exist_ok=True)
  path.write_text(json.dumps(document, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')


def slugify(value: str) -> str:
  return re.sub(r'[^a-z0-9]+', '-', value.strip().lower()).strip('-') or 'score'


def human_label(value: str) -> str:
  return ' '.join(v.capitalize() for v in value.replace('_',' ').replace('-',' ').split()) or 'General'


def source_staff(document: dict) -> dict:
  if len(document.get('staves', [])) != 1:
    raise ValueError('exactly one source staff is required')
  return document['staves'][0]


def first_voice(measure: dict) -> dict:
  if len(measure.get('voices', [])) != 1:
    raise ValueError('exactly one source voice is required')
  return measure['voices'][0]


def measure_notes(measure: dict) -> list[dict]:
  return first_voice(measure)['notes']


def leading_pickup_measure_index(measures: list[dict]) -> int | None:
  return 0 if len(measures)>1 and measures[0].get('implicit') else None


def body_start_measure_index(measures: list[dict]) -> int:
  return 1 if leading_pickup_measure_index(measures) is not None else 0


def copy_source_voice(measure: dict, *, strip_lyrics: bool) -> dict:
  voice = copy.deepcopy(first_voice(measure))
  if strip_lyrics:
    for note in voice['notes']:
      note.pop('lyric', None)
      note.pop('lyric_notation', None)
  return voice


def make_measure(measure_id: str, source: dict, notes: list[dict], *, section_label: str = '') -> dict:
  result = copy.deepcopy(source)
  result.update({'id':measure_id, 'section_label':section_label, 'voices':[{'id':'voice-arranged', 'notes':notes}]})
  return result


def reidentify(measures: list[dict], source_tuplets: list[dict]) -> list[dict]:
  """Clone-local links survive; links cut by an arrangement boundary become attacks."""
  tone_map, event_map = {}, {}
  tuplets = []
  for measure in measures:
    for ni, note in enumerate(measure_notes(measure)):
      old = note['id']
      note['id'] = f'{measure["id"]}-n{ni+1}'
      event_map[old] = note['id']
      for pi, pitch in enumerate(note['pitches']):
        old_pitch = pitch['id']
        pitch['id'] = f'{note["id"]}-p{pi+1}'
        tone_map[old_pitch] = pitch['id']
  for measure in measures:
    for note in measure_notes(measure):
      for p in note['pitches']:
        for field in ('tie_from','tie_to'):
          if p.get(field) in tone_map:
            p[field] = tone_map[p[field]]
          else:
            p.pop(field,None)
      if note.get('tremolo', {}).get('partner'):
        partner = note['tremolo']['partner']
        if partner in event_map:
          note['tremolo']['partner'] = event_map[partner]
        else:
          note['tremolo'] = {**note['tremolo'], 'type':'single'}
          note['tremolo'].pop('partner',None)
  for source in source_tuplets:
    members = [n for m in measures for n in measure_notes(m) if source['id'] in n.get('tuplet_ids', [])]
    if not members:
      continue
    t = copy.deepcopy(source)
    t.update({'id':f'{measures[0]["id"]}-{source["id"]}', 'start':members[0]['id'], 'end':members[-1]['id'], 'continued_start':source['start'] not in event_map, 'continued_end':source['end'] not in event_map})
    for n in members:
      n['tuplet_ids'] = [t['id'] if v==source['id'] else v for v in n['tuplet_ids']]
    tuplets.append(t)
  return tuplets


def make_arranged_document(source_document: dict, staff: dict, *, staff_id: str, staff_name: str, role: str, arrangement_part_type: str, measures: list[dict], tuplets=None, octave_shifts=None) -> dict:
  return {'schema_version':SCHEMA_VERSION, 'metadata':copy.deepcopy(source_document['metadata']), 'staves':[{'id':staff_id, 'name':staff_name, 'clef':copy.deepcopy(staff['clef']), 'role':role, 'arranged':True, 'arrangement_part_type':arrangement_part_type, 'measures':measures, 'tuplets':tuplets or [], 'octave_shifts':octave_shifts or []}]}


def count_json_files(root: Path) -> int:
  return sum(1 for p in root.rglob('*.json') if p.is_file())


def validate_generated_tree(root: Path) -> None:
  for path in root.rglob('*.json'):
    load_json(path)