from __future__ import annotations

from pathlib import Path

from generator.arrange.form import ArrangementSection, full_arrangement_sections
from generator.arrange.fretboard import pitches_for_voicing
from generator.score import (
  QUARTER_TICKS,
  apply_source_chord_markers,
  active_chord_at_offset,
  chord_markers,
  chord_symbol_at_offset,
  copy_source_voice,
  first_chord_symbol,
  first_voice,
  leading_pickup_measure_index,
  load_json,
  make_arranged_document,
  make_measure,
  measure_capacity_ticks,
  renderable_duration,
  renderable_rest_notes,
  rest_note,
  sanitized_note,
  slash_note,
  slugify,
  source_staff,
  source_note_attack_at,
  write_json,
)

PART_TYPES = ("acoustic", "electric", "vocal")


def propagated_chords(measures: list[dict]) -> list[str]:
  result: list[str] = []
  current = ""
  for measure in measures:
    current = first_chord_symbol(measure, current)
    result.append(current)
  return result


def playable_chord_note(note_id: str, duration_ticks: int, voicing_chord_symbol: str, visible_chord_symbol: str = "") -> dict:
  note = {
    "id": note_id,
    "kind": "note",
    "duration_ticks": duration_ticks,
    "pitches": pitches_for_voicing(voicing_chord_symbol),
    "hidden": False,
  }
  if not note["pitches"]:
    return rest_note(note_id, duration_ticks, visible_chord_symbol)
  if visible_chord_symbol:
    note["chord_symbol"] = visible_chord_symbol
  return note


def rhythm_notes(measure: dict, chord_symbol: str, sequence: int) -> tuple[list[dict], int]:
  capacity = measure_capacity_ticks(str(measure.get("time_signature", "4/4")))
  hit_offsets = (0, QUARTER_TICKS + QUARTER_TICKS // 2, QUARTER_TICKS * 2, QUARTER_TICKS * 2 + QUARTER_TICKS // 2, QUARTER_TICKS * 3 + QUARTER_TICKS // 2)
  hit_duration = QUARTER_TICKS // 2
  notes: list[dict] = []
  cursor = 0
  for offset in hit_offsets:
    if offset >= capacity:
      break
    if offset > cursor:
      rests, sequence = renderable_rest_notes("arr-rest", offset - cursor, sequence)
      notes.extend(rests)
      cursor = offset
    active_chord = active_chord_at_offset(measure, offset, chord_symbol)
    visible_chord = chord_symbol_at_offset(measure, offset)
    sequence += 1
    notes.append(playable_chord_note(f"arr-note-{sequence}", min(hit_duration, capacity - offset), active_chord, visible_chord))
    cursor = offset + min(hit_duration, capacity - offset)
  if cursor < capacity:
    rests, sequence = renderable_rest_notes("arr-rest", capacity - cursor, sequence)
    notes.extend(rests)
  return apply_source_chord_markers(notes, measure), sequence


def accompaniment_notes(measure: dict, chord_symbol: str, sequence: int) -> tuple[list[dict], int]:
  capacity = measure_capacity_ticks(str(measure.get("time_signature", "4/4")))
  notes: list[dict] = []
  lead_rest = min(QUARTER_TICKS // 2, capacity)
  if lead_rest:
    rests, sequence = renderable_rest_notes("arr-rest", lead_rest, sequence, chord_symbol_at_offset(measure, 0))
    notes.extend(rests)
  if lead_rest < capacity:
    hit = min(QUARTER_TICKS // 2, capacity - lead_rest)
    visible_chord = chord_symbol_at_offset(measure, lead_rest)
    active_chord = active_chord_at_offset(measure, lead_rest, chord_symbol)
    sequence += 1
    if source_note_attack_at(measure, lead_rest):
      notes.append(rest_note(f"arr-rest-{sequence}", hit, visible_chord))
    else:
      notes.append(playable_chord_note(f"arr-note-{sequence}", hit, active_chord, visible_chord))
    cursor = lead_rest + hit
    if cursor < capacity:
      markers = [marker for marker in chord_markers(measure) if marker[0] >= cursor]
      if markers:
        segment_cursor = cursor
        for index, (start, _, chord) in enumerate(markers):
          if start > segment_cursor:
            rests, sequence = renderable_rest_notes("arr-rest", start - segment_cursor, sequence)
            notes.extend(rests)
          next_start = markers[index + 1][0] if index + 1 < len(markers) else capacity
          rests, sequence = renderable_rest_notes("arr-rest", next_start - start, sequence, chord)
          notes.extend(rests)
          segment_cursor = next_start
        if segment_cursor < capacity:
          rests, sequence = renderable_rest_notes("arr-rest", capacity - segment_cursor, sequence)
          notes.extend(rests)
      else:
        rests, sequence = renderable_rest_notes("arr-rest", capacity - cursor, sequence)
        notes.extend(rests)
  return apply_source_chord_markers(notes, measure), sequence


def slash_notes(measure: dict, chord_symbol: str, sequence: int) -> tuple[list[dict], int]:
  capacity = measure_capacity_ticks(str(measure.get("time_signature", "4/4")))
  notes: list[dict] = []
  cursor = 0
  while cursor < capacity:
    sequence += 1
    duration = min(QUARTER_TICKS, capacity - cursor)
    notes.append(slash_note(f"arr-slash-{sequence}", duration))
    cursor += duration
  return notes, sequence


def rest_notes(measure: dict, chord_symbol: str, sequence: int) -> tuple[list[dict], int]:
  capacity = measure_capacity_ticks(str(measure.get("time_signature", "4/4")))
  markers = chord_markers(measure)
  if not markers:
    notes, sequence = renderable_rest_notes("arr-rest", capacity, sequence)
    return notes, sequence

  notes: list[dict] = []
  cursor = 0
  for index, (start, _, chord) in enumerate(markers):
    if start > cursor:
      if not renderable_duration(start - cursor):
        return source_rest_skeleton_notes(measure, sequence)
      rests, sequence = renderable_rest_notes("arr-rest", start - cursor, sequence)
      notes.extend(rests)
    next_start = markers[index + 1][0] if index + 1 < len(markers) else capacity
    if not renderable_duration(next_start - start):
      return source_rest_skeleton_notes(measure, sequence)
    rests, sequence = renderable_rest_notes("arr-rest", next_start - start, sequence, chord)
    notes.extend(rests)
    cursor = next_start
  if cursor < capacity:
    if not renderable_duration(capacity - cursor):
      return source_rest_skeleton_notes(measure, sequence)
    rests, sequence = renderable_rest_notes("arr-rest", capacity - cursor, sequence)
    notes.extend(rests)
  return notes, sequence


def source_rest_skeleton_notes(measure: dict, sequence: int) -> tuple[list[dict], int]:
  notes: list[dict] = []
  for source_note in first_voice(measure).get("notes", []):
    if not isinstance(source_note, dict):
      continue
    sequence += 1
    note = rest_note(f"arr-rest-{sequence}", int(source_note.get("duration_ticks", QUARTER_TICKS)), str(source_note.get("chord_symbol", "")))
    if "display_ticks" in source_note:
      note["display_ticks"] = source_note["display_ticks"]
    if "tuplet_id" in source_note:
      note["tuplet_id"] = source_note["tuplet_id"]
    notes.append(note)
  if not notes:
    sequence += 1
    notes.append(rest_note(f"arr-rest-{sequence}", measure_capacity_ticks(str(measure.get("time_signature", "4/4")))))
  return notes, sequence


def melody_voice(measure: dict, *, strip_lyrics: bool) -> dict:
  voice = copy_source_voice(measure, strip_lyrics=strip_lyrics)
  voice["id"] = "voice-arranged-source"
  voice["notes"] = [sanitized_note(note, strip_lyrics=strip_lyrics) for note in first_voice(measure).get("notes", []) if isinstance(note, dict)]
  return voice


def apply_vocal_pickup_overlap(output: list[dict], source_measures: list[dict], section: ArrangementSection) -> None:
  pickup = leading_pickup_measure_index(source_measures)
  if pickup is None or section.source_measure_index != pickup + 1 or not output:
    return
  voice = copy_source_voice(source_measures[pickup], strip_lyrics=False)
  for note in voice["notes"]:
    if isinstance(note, dict):
      note.pop("chord_symbol", None)
  output[-1]["voices"][0]["id"] = "voice-arranged-vocal"
  output[-1]["voices"][0]["tuplets"] = voice["tuplets"]
  output[-1]["voices"][0]["notes"] = voice["notes"]


def voice_for_section(section: ArrangementSection, measure: dict, chord_symbol: str, sequence: int) -> tuple[dict, int]:
  if section.role == "vocal":
    return melody_voice(measure, strip_lyrics=False), sequence
  if section.role == "melody":
    return melody_voice(measure, strip_lyrics=True), sequence
  if section.role == "rhythm":
    notes, sequence = rhythm_notes(measure, chord_symbol, sequence)
  elif section.role == "accompaniment":
    notes, sequence = accompaniment_notes(measure, chord_symbol, sequence)
  elif section.role == "slash":
    notes, sequence = slash_notes(measure, chord_symbol, sequence)
  else:
    notes, sequence = rest_notes(measure, chord_symbol, sequence)
  return {"id": "voice-arranged-guitar", "tuplets": [], "notes": notes}, sequence


def arranged_measures(source_measures: list[dict], part_type: str) -> list[dict]:
  sections = full_arrangement_sections(source_measures, part_type)
  chords = propagated_chords(source_measures)
  output: list[dict] = []
  sequence = 0
  for section in sections:
    if part_type == "vocal" and section.role == "vocal":
      apply_vocal_pickup_overlap(output, source_measures, section)
    for offset in range(section.measure_count):
      source_index = section.source_measure_index + offset
      source_measure = source_measures[source_index]
      chord_symbol = chords[source_index] if source_index < len(chords) else ""
      voice, sequence = voice_for_section(section, source_measure, chord_symbol, sequence)
      output_measure = make_measure(
        f"measure-{len(output) + 1}",
        source_measure,
        voice.get("notes", []),
        section_label=section.label if offset == 0 else "",
        tuplets=voice.get("tuplets", []),
      )
      output_measure["voices"][0]["id"] = str(voice.get("id", output_measure["voices"][0]["id"]))
      output.append(output_measure)
  return output


def arrange_document(document: dict, part_type: str) -> dict:
  staff = source_staff(document)
  measures = staff.get("measures")
  if not isinstance(measures, list) or not measures:
    raise ValueError("performance arrangement requires at least one source measure")
  output = make_arranged_document(
    document,
    staff,
    staff_id="staff-arranged-vocal" if part_type == "vocal" else "staff-arranged-guitar",
    staff_name=part_type.capitalize(),
    role="vocal" if part_type == "vocal" else "guitar",
    arrangement_part_type=part_type,
    measures=arranged_measures([measure for measure in measures if isinstance(measure, dict)], part_type),
  )
  metadata = output.setdefault("metadata", {})
  if isinstance(metadata, dict):
    metadata.setdefault("tempo", 124)
  return output


def generate_performance_scores(source_dir: Path, output_dir: Path) -> int:
  if not source_dir.exists():
    raise ValueError(f"performance source directory not found: {source_dir}")
  count = 0
  for source_path in sorted(source_dir.glob("*.json")):
    document = load_json(source_path)
    for part_type in PART_TYPES:
      output = arrange_document(document, part_type)
      target = output_dir / f"{slugify(source_path.stem)}-{part_type}.json"
      write_json(target, output)
      count += 1
  return count