from __future__ import annotations

import copy
import hashlib
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'))
  if 'music' in document:
    from generator.engraving import ENGINE_VERSION, PAGE_HEIGHT, PAGE_WIDTH, check_page
    presentation = document.get('presentation', {})
    if document.get('schema_version') != SCHEMA_VERSION or not isinstance(presentation, dict) or not isinstance(presentation.get('pages'), list) or not presentation['pages']:
      raise ValueError('invalid score bundle; regenerate score assets')
    if (presentation.get('engine'), presentation.get('engine_version'), presentation.get('width'), presentation.get('height')) != ('MuseScore Studio', ENGINE_VERSION, PAGE_WIDTH, PAGE_HEIGHT):
      raise ValueError('unsupported score presentation; regenerate score assets')
    for page in presentation['pages']:
      if not isinstance(page, dict) or not isinstance(page.get('path'), str) or not isinstance(page.get('sha256'), str):
        raise ValueError('invalid score page manifest')
      name = page['path']
      relative = Path(name)
      if '\x00' in name or '\\' in name or relative.is_absolute() or any(part in ('', '.', '..') for part in name.split('/')) or relative.suffix != '.png':
        raise ValueError('invalid score page path')
      resolved = (path.parent / relative).resolve()
      if not resolved.is_relative_to(path.parent.resolve()):
        raise ValueError('score page escapes its asset directory')
      if not re.fullmatch('[0-9a-f]{64}', page['sha256']) or hashlib.sha256(resolved.read_bytes()).hexdigest() != page['sha256']:
        raise ValueError('score page checksum mismatch')
      check_page(resolved)
    document = {'schema_version': SCHEMA_VERSION, **document['music']}
  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 measure_notes(measure: dict) -> list[dict]:
  return [note for voice in measure['voices'] for note in voice['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_voices(measure: dict, *, strip_lyrics: bool) -> list[dict]:
  voices = copy.deepcopy(measure['voices'])
  if strip_lyrics:
    for note in (n for v in voices for n in v['notes']):
      note.pop('lyric', None)
      note.pop('lyric_notation', None)
      note.pop('lyrics', None)
  return voices

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}]})
  for note in notes:
    note['voice'] = 'voice-arranged'
  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 hi, harmony in enumerate(measure.get('harmonies', [])):
      harmony['id'] = f'{measure["id"]}-h{hi + 1}'
    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 [], 'instruments': copy.deepcopy(staff.get('instruments', []))}]}

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)