from __future__ import annotations

from pathlib import Path

from generator.score import copy_source_voice, make_arranged_document, make_measure, reidentify, slugify, source_staff, write_json
from generator.musicxml import load_musicxml
from generator.transposition import TARGET_KEYS, transpose_document


def practice_document(document: dict) -> dict:
  staff = source_staff(document)
  source_measures = staff.get("measures")
  if not isinstance(source_measures, list):
    raise ValueError("practice source must contain measures")
  measures: list[dict] = []
  for index, measure in enumerate(source_measure for source_measure in source_measures if isinstance(source_measure, dict)):
    voice = copy_source_voice(measure, strip_lyrics=True)
    output_measure = make_measure(
      f"measure-{index + 1}",
      measure,
      voice["notes"],
      section_label=str(measure.get("section_label", "")),
    )
    output_measure["voices"][0]["id"] = str(voice.get("id", output_measure["voices"][0]["id"]))
    measures.append(output_measure)
  tuplets = reidentify(measures, staff.get('tuplets', []))
  return make_arranged_document(
    document,
    staff,
    staff_id="staff-arranged-practice",
    staff_name="Practice",
    role="vocal",
    arrangement_part_type="practice",
    measures=measures,
    tuplets=tuplets,
    octave_shifts=staff.get('octave_shifts', []),
  )


def key_file_name(key: str) -> str:
  return slugify(key.replace("#", "sharp")) + ".json"


def generate_chord_practice(source_path: Path, source_root: Path, output_dir: Path) -> int:
  document = practice_document(load_musicxml(source_path))
  metadata = document.get("metadata", {})
  key = str(metadata.get("key_signature", "C") if isinstance(metadata, dict) else "C")
  relative = source_path.relative_to(source_root)
  lesson_group = relative.parent.name
  exercise_id = slugify(source_path.stem)
  write_json(output_dir / "chord" / lesson_group / exercise_id / key_file_name(key), document)
  return 1


def generate_melody_practice(source_path: Path, source_root: Path, output_dir: Path) -> int:
  base_document = practice_document(load_musicxml(source_path))
  relative = source_path.relative_to(source_root)
  lesson_group = relative.parent.name
  exercise_id = slugify(source_path.stem)
  count = 0
  for key in TARGET_KEYS:
    write_json(output_dir / "melody" / lesson_group / exercise_id / key_file_name(key), transpose_document(base_document, key))
    count += 1
  return count


def generate_practice_scores(source_dir: Path, output_dir: Path) -> int:
  if not source_dir.exists():
    return 0
  count = 0
  chord_root = source_dir / "chord"
  if chord_root.exists():
    for path in sorted(chord_root.rglob("*.musicxml")):
      count += generate_chord_practice(path, chord_root, output_dir)
  melody_root = source_dir / "melody"
  if melody_root.exists():
    for path in sorted(melody_root.rglob("*.musicxml")):
      count += generate_melody_practice(path, melody_root, output_dir)
  return count