"""Independent, reproducible role generation and non-composing assembly.

Only source documents and production resources enter generate_role. Completed
roles and published assets are downstream values, never composition inputs.
"""
from __future__ import annotations
import concurrent.futures
import copy
import dataclasses
import hashlib
import functools
import json
import multiprocessing
import os
import pathlib
import tempfile
import time
import fractions

import generator.accompaniment
import generator.comping
import generator.fretboard
import generator.model
import generator.musicxml
import generator.parts
import generator.score
import generator.vocabulary
import generator.arrange.accompaniment
import generator.arrange.ending
import generator.arrange.form
import generator.arrange.generate
import generator.arrange.pickup
import generator.solo.analysis
import generator.solo.composer
import generator.solo.notation
import generator.solo.report
import generator.solo.learned
import generator.solo.authored
import generator.input_audit
import generator.solo.arrangement

VERSION = 'independent-roles-1'
ROLES = ('acoustic_accompaniment', 'electric_accompaniment', 'electric_melody', 'vocal_melody', 'acoustic_solo', 'electric_solo')

@dataclasses.dataclass(frozen=True)
class SourceAnalysis:
  document_json: str
  musical_key: str
  @classmethod
  def read(cls, document):
    prepared = copy.deepcopy(document)
    staff = generator.score.source_staff(prepared)
    for i, m in enumerate(staff['measures']):
      m['id'], m['number'] = f'source_{i}', str(i + 1)
    staff['tuplets'] = generator.score.reidentify(staff['measures'], staff.get('tuplets', []))
    prepared['metadata'] = {k: v for k, v in prepared['metadata'].items() if k not in ('title', 'composer', 'copyright', 'software')}
    prepared['metadata'].update(title='Source', composer='')
    staff.update(id='source', name='Source')
    melodic = generator.solo.analysis.analyze(prepared)
    content = {'melody': [(str(n.start), str(n.end), dataclasses.asdict(n.pitch)) for n in melodic.melody],
      'measures': [{k: m.get(k) for k in ('duration', 'time_signature', 'key', 'harmonies', 'tempos', 'section_label', 'implicit')} for m in staff['measures']]}
    # Fraction/string spelling is canonical and never includes filesystem IDs.
    key = hashlib.sha256(json.dumps(content, sort_keys=True, default=str).encode()).hexdigest()
    return cls(json.dumps(prepared, sort_keys=True), key)
  def document(self):
    return json.loads(self.document_json)
  def accompaniment_key(self):
    document = self.document()
    document['metadata'].pop('tempo', None)
    for measure in generator.score.source_staff(document)['measures']:
      measure['tempos'] = []
    return SourceAnalysis.read(document).musical_key

def windows(document, part):
  source = generator.score.source_staff(document)['measures']
  sections = generator.arrange.form.full_arrangement_sections(source, part)
  result, cursor = [], fractions.Fraction()
  for section in sections:
    duration = sum((generator.model.fraction(m['duration']) for m in source[section.source_measure_index:section.source_measure_index + section.measure_count]), fractions.Fraction())
    role = {'slash': 'solo', 'vocal': 'melody'}.get(section.role, section.role)
    if section.role in ('vocal', 'melody') and generator.arrange.pickup.enters(section, source) and result:
      attacks = [generator.model.fraction(n['onset']) for n in generator.score.measure_notes(source[0]) if n['pitches']]
      if attacks:
        start = cursor - generator.model.fraction(source[0]['duration']) + min(attacks)
        a, _, old = result[-1]
        result[-1] = (a, start, old)
        result.append((start, cursor, part + '_' + role))
    end = cursor + duration
    start = cursor
    if section.role == 'slash' and generator.score.body_start_measure_index(source) and result:
      # Allocate the source's complete pickup bar before composition. This
      # does not inspect either generated solo; a role may leave it silent.
      start = cursor - generator.model.fraction(source[0]['duration'])
      a, _, old = result[-1]
      result[-1] = (a, start, old)
    result.append((start, end, part + '_' + role))
    cursor = end
  result.append((cursor, cursor + generator.arrange.ending.duration(source), part + ('_accompaniment' if part != 'vocal' else '_rest')))
  return tuple(result)

def dependencies(role, resources):
  names = ['harmony', 'lead_ins'] if role.endswith('_solo') else []
  return {'roles': {'role': resources['roles']['roles'][role], 'difficulty': resources['roles']['difficulty']},
    **({'voicing_families': generator.vocabulary.accompaniment('electric', resources)['voicing_families']} if role == 'electric_melody' else {}),
    **({'accompaniment': generator.vocabulary.accompaniment(role.split('_')[0], resources)} if role.endswith('_accompaniment') else {}),
    **{name: resources[name] for name in names}}

def _code_key(role):
  root = pathlib.Path(__file__).parent
  paths = [root / n for n in ('production.py', 'model.py', 'harmony.py', 'pitch.py', 'fretboard.py', 'parts.py', 'score.py', 'vocabulary.py')]
  paths += list((root / 'arrange').glob('*.py'))
  if role.endswith('_solo'):
    paths += list((root / 'solo').glob('*.py'))
    paths += [root / 'harmonic_time.py', root / 'musicxml.py', root / 'training_io.py']
  elif role.endswith('_accompaniment'):
    paths += [root / n for n in ('comping.py', 'comping_rhythm.py', 'chord_families.py', 'accompaniment.py', 'guitar.py')]
    if role == 'acoustic_accompaniment':
      paths += [root / 'acoustic_voicing.py', root / 'acoustic.py']
  elif role == 'electric_melody':
    paths += [root / 'guitar.py', root / 'melody.py', root / 'chord_families.py']
  return generator.vocabulary.fingerprint([(p.relative_to(root).as_posix(), hashlib.sha256(p.read_bytes()).hexdigest()) for p in sorted(set(paths))])

def _resource_key(directory, role):
  return {'model': hashlib.sha256((directory / 'models' / 'solo.json').read_bytes()).hexdigest()} if role.endswith('_solo') else {}

@functools.lru_cache(maxsize=4)
def _solo_model(directory, fingerprint):
  return generator.solo.learned.Model.load(pathlib.Path(directory) / 'models' / 'solo.json')

def _mask(document, allowed):
  """Apply source-defined role exclusions, without observing another part."""
  absolute = fractions.Fraction()
  for measure in document['staves'][0]['measures']:
    length = generator.model.fraction(measure['duration'])
    for voice in measure['voices']:
      selected = []
      for n in voice['notes']:
        start = absolute + generator.model.fraction(n['onset'])
        end = start + generator.model.fraction(n['duration'])
        intervals = [(max(start, a), min(end, b)) for a, b in allowed if a < end and b > start]
        for a, b in intervals:
          note = copy.deepcopy(n)
          if a != start:
            note['onset'] = generator.model.rational(a - absolute)
            for pitch in note['pitches']:
              pitch.pop('tie_from', None)
          generator.arrange.pickup.truncate(note, b - absolute)
          if a < b:
            selected.append(note)
      voice['notes'] = _fill(selected, length, measure['id'] + '_' + voice['id'], voice['id'])
    absolute += length
  staff = document['staves'][0]
  staff['tuplets'] = generator.arrange.pickup.retain_links(staff['measures'], staff.get('tuplets', []))

def _fill(notes, length, prefix, voice='voice_arranged'):
  notes = sorted(notes, key=lambda n: generator.model.fraction(n['onset']))
  result = []
  cursor = fractions.Fraction()

  def rests(end):
    nonlocal cursor
    for duration, written in generator.accompaniment.written_durations(end - cursor):
      result.append(generator.model.make_event(prefix + '_rest_' + str(len(result)), cursor, duration, 'rest', written=written) | {'voice': voice})
      cursor += duration

  for n in notes:
    start = generator.model.fraction(n['onset'])
    if start < cursor:
      raise ValueError(prefix + ': overlapping independently generated roles')
    rests(start)
    n['voice'] = voice
    result.append(n)
    cursor = start + generator.model.fraction(n['duration'])
  rests(length)
  return result

def generate_role(source, role, data_dir, settings=None, cache_dir=None):
  started = time.monotonic()
  with generator.input_audit.composition(data_dir, role, cache_dir) as reads:
    result = _generate_role(source, role, data_dir, settings, cache_dir)
  # Observed reads are diagnostics, never cache identity or musical seeds.
  result['trace'] = {**result['trace'], 'input_boundary': 'production data (including explicit authored repertoire), generator code, Python runtime, own role cache'}
  result['observed_reads'] = sorted(reads)
  result['seconds'] = time.monotonic() - started
  return result

def _generate_role(source, role, data_dir, settings=None, cache_dir=None):
  if role not in ROLES:
    raise ValueError('Unknown generation role: ' + role)
  data_dir = pathlib.Path(data_dir)
  resources = generator.vocabulary.load(data_dir / 'composition')
  settings = settings or {}
  registered = generator.solo.authored.lookup(source, role.removesuffix('_solo'), data_dir) if role.endswith('_solo') and not settings.get('predict', False) else None
  algorithm = VERSION + '/' + generator.solo.learned.VERSION if role.endswith('_solo') else VERSION
  inputs = {'version': algorithm, 'source': source.musical_key, 'role': role, 'settings': settings,
    'vocabulary': dependencies(role, resources), 'catalogs': ({'authored': registered['sha256'], 'id': registered['id'], 'pitch_convention': registered['pitch_convention']} if registered else _resource_key(data_dir, role)), 'code': _code_key(role)}
  # Source notation/lyrics also invalidate cached score representation, never the musical seed.
  key = generator.vocabulary.fingerprint({'inputs': inputs, 'notation': source.document_json})
  path = pathlib.Path(cache_dir) / role / (key + '.json') if cache_dir else None
  if path and path.is_file() and not path.is_symlink():
    try:
      entry = json.loads(path.read_text())
      value = entry['result']
      if entry['key'] == key and entry['digest'] == generator.vocabulary.fingerprint(value):
        generator.model.validate_document(value['document'])
        generator.fretboard.validate_document(value['document'])
        return {**value, 'cache_hit': True}
    except (ValueError, KeyError, TypeError):
      pass
  started = time.monotonic()
  document = source.document()
  part = role.split('_')[0]
  trace = {'source': source.musical_key, 'role': role, 'production_resources': inputs['catalogs'],
    'vocabulary': generator.vocabulary.fingerprint(inputs['vocabulary']), 'settings': settings, 'algorithm': algorithm, 'code': inputs['code'], 'cache_key': key}
  choice_source = source.accompaniment_key() if role.endswith('_accompaniment') else source.musical_key
  choice_data, choice_settings = inputs['vocabulary'], settings
  if role.endswith('_accompaniment'):
    # Voicing weights invalidate the concrete cache, but must not reshuffle
    # rhythm patterns. Removing this new field preserves the pre-policy seed.
    choice_data = copy.deepcopy(choice_data)
    choice_data['roles']['role'].pop('voicing', None)
    choice_settings = {k: v for k, v in settings.items() if k != 'voicing'}
  seed = generator.vocabulary.fingerprint({'source': choice_source, 'role': role, 'settings': choice_settings, 'data': choice_data, 'version': algorithm})
  details = {}
  solo_report = ''
  composition_timings = {}
  with generator.vocabulary.using(resources, role, settings | {'choice_seed': seed}):
    if role.endswith('_solo'):
      result = generator.arrange.generate._arrange_document(document, part, {}, owned={'slash'})
      si = 3 if part == 'acoustic' else 4
      first, count, offset, length = generator.solo.arrangement.regions(document)[si]
      # The template/harmony comes directly from source form, not another generated part.
      authored = generator.score.source_staff(document)['measures']
      section = generator.arrange.form.full_arrangement_sections(authored, part)[si]
      effective = generator.arrange.generate.section_source(authored, section, generator.arrange.form.full_arrangement_sections(authored, part)[si + 1])
      template = copy.deepcopy(effective[section.source_measure_index:section.source_measure_index + count])
      initial = generator.accompaniment.carried_chords(effective)[section.source_measure_index]
      start, end = next((a - offset, b - offset) for a, b, r in windows(document, part) if r == role)
      notation_first = first
      if start < 0:
        # The preceding arrangement harmony comes from the source form. Do
        # not invent harmony from a comparison score's unlabelled lead-in.
        template.insert(0, copy.deepcopy(result['staves'][0]['measures'][first - 1]))
        notation_first -= 1
        initial = generator.accompaniment.carried_chords(result['staves'][0]['measures'])[first - 1]
      if registered:
        solo = generator.solo.authored.realize(registered, generator.solo.analysis.analyze(document), template, part, start, end, initial)
      else:
        solo_model = _solo_model(str(data_dir), inputs['catalogs']['model'])
        solo = generator.solo.composer.compose(generator.solo.analysis.analyze(document), template, part, end, initial=initial, learned_model=solo_model, expressions_enabled=generator.vocabulary.role().get('expressions', True), start=start)
      composition_timings = {k: v for k, v in solo.diagnostics.items() if k.endswith('_seconds')}
      relative = tuple(dataclasses.replace(n, start=n.start - start, end=n.end - start) for n in solo.notes())
      output = generator.solo.notation.measures(result['staves'][0]['measures'][notation_first:first + count], relative, 'solo_' + part)
      generator.solo.notation.replace_region(result, notation_first, output)
      details = generator.solo.report.structure(solo)
      details['allocation'] = {'start': str(start), 'end': str(end), 'lead_in_measures': first - notation_first}
      solo_report = generator.solo.report.describe(solo, first, offset, detail=details)
    elif role.endswith('_accompaniment'):
      canonical = generator.comping.Canonical(document, part, seed)
      result = generator.arrange.generate._arrange_document(document, part, {}, owned={'accompaniment'}, accompaniment=canonical)
      details = {'canonical_comping': canonical.report, 'accompaniment_rhythm': canonical.rhythm_report, 'groove_cycles': canonical.groove_report, 'rhythm_source': choice_source}
    else:
      result = generator.arrange.generate._arrange_document(document, part, {}, owned={'vocal' if part == 'vocal' else 'melody'})
  allowed = tuple((a, b) for a, b, r in windows(document, part) if r == role)
  _mask(result, allowed)
  if role == "acoustic_accompaniment":
    generator.acoustic.validate_comping(result)
  staff = result['staves'][0]
  for measure in staff['measures']:
    measure['id'] = role + '_' + measure['id']
  staff['tuplets'] = generator.score.reidentify(staff['measures'], staff.get('tuplets', []))
  if part != 'vocal':
    result['staves'][0]['tablature'] = True
  generator.model.validate_document(result)
  generator.fretboard.validate_document(result)
  value = {'role': role, 'document': result, 'trace': trace, 'details': details, 'report': solo_report, 'seconds': time.monotonic() - started, 'composition_timings': composition_timings, 'cache_hit': False}
  # Sidecar diagnostics use the same rational-safe representation as reports.
  # Identical insertion order after cold generation and cache loading also
  # keeps serialized scores and the catalog byte-for-byte reproducible.
  value = json.loads(json.dumps(value, default=str, sort_keys=True))
  if path:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {'key': key, 'digest': generator.vocabulary.fingerprint(value), 'result': value}
    with tempfile.NamedTemporaryFile('w', dir=path.parent, delete=False) as out:
      json.dump(payload, out, sort_keys=True)
      temporary = out.name
    os.replace(temporary, path)
  return value

def assemble(source, results):
  """Union prevalidated role windows; no selection or pitch/fingering rewrite."""
  documents = {}
  for part in generator.parts.PART_TYPES:
    selected = [results[r]['document'] for r in ROLES if r.startswith(part + '_')]
    document = copy.deepcopy(selected[0])
    staff = document['staves'][0]
    if part == 'vocal':
      documents[part] = document
      continue
    staff['tuplets'] = []
    staff['octave_shifts'] = []
    staff['scale_spans'] = []
    for other in selected:
      for field in ('tuplets', 'octave_shifts', 'scale_spans'):
        staff[field].extend(copy.deepcopy(other['staves'][0].get(field, [])))
    for i, m in enumerate(staff['measures']):
      # Tuplet rests belong to the source/role rhythm and must survive union.
      lanes = {}
      for d in selected:
        for voice in d['staves'][0]['measures'][i]['voices']:
          lane = voice['id'] if part == 'acoustic' and voice['id'] in ('upper', 'bass') else 'voice_arranged'
          sounding = [copy.deepcopy(n) for n in voice['notes'] if n['pitches'] or n.get('tuplet_ids')]
          if sounding:
            lanes.setdefault(lane, []).extend(sounding)
      lanes = lanes or {'voice_arranged': []}
      m['voices'] = [{'id': lane, 'notes': _fill(notes, generator.model.fraction(m['duration']), m['id'] + '_' + lane, lane)} for lane, notes in lanes.items()]
      # Attribute changes belong to source melody pickups and are assembly metadata.
      for d in selected:
        other = d['staves'][0]['measures'][i]
        if other.get('attributes'):
          m['attributes'] = copy.deepcopy(other['attributes'])
    staff['tuplets'] = generator.arrange.pickup.retain_links(staff['measures'], staff['tuplets'])
    owners = {n['id']: v['id'] for m in staff['measures'] for v in m['voices'] for n in v['notes']}
    for span in staff['tuplets']:
      span['voice'] = owners[span['start']]
    generator.model.validate_document(document)
    generator.fretboard.validate_document(document)
    documents[part] = document
  return documents

def _job(request):
  source, role, directory, settings, cache = request
  return generate_role(source, role, directory, settings, cache)

def generate_tree(source_dir, output_dir, *, jobs=2, cache_dir=None, progress=False, timings=None):
  if type(jobs) is not int or not 1 <= jobs <= 16:
    raise ValueError('composition workers must be between 1 and 16')
  source_dir, output_dir = pathlib.Path(source_dir), pathlib.Path(output_dir)
  sources = {}
  for path in generator.parts.source_paths(source_dir):
    document = generator.parts.prepare_document(generator.musicxml.load_musicxml(path))
    sources[path] = (document, SourceAnalysis.read(document))
  requests = [(path, role, (s, role, source_dir.parent, {}, cache_dir)) for path, (_, s) in sources.items() for role in ROLES]
  results = {p: {} for p in sources}

  def accept(path, role, value):
    results[path][role] = value
    if timings is not None:
      timings.append({'source': str(path), 'role': role, 'seconds': value['seconds'], 'cache_hit': value['cache_hit'], 'composition_timings': value.get('composition_timings', {}), 'observed_reads': value.get('observed_reads', [])})
    if progress:
      print(f'Composition: {path.stem}/{role}: ' + ('cached' if value['cache_hit'] else f'{value["seconds"]:.1f}s'), flush=True)

  if jobs == 1:
    for path, role, args in requests:
      try:
        accept(path, role, _job(args))
      except Exception as error:
        raise ValueError(f'{path.name}/{role}: {error}') from error
  else:
    with concurrent.futures.ProcessPoolExecutor(max_workers=jobs, mp_context=multiprocessing.get_context('spawn')) as pool:
      futures = {pool.submit(_job, args): (path, role) for path, role, args in requests}
      for future in concurrent.futures.as_completed(futures):
        path, role = futures[future]
        try:
          accept(path, role, future.result())
        except Exception as error:
          raise ValueError(f'{path.name}/{role}: {error}') from error
  for path, (original, source) in sources.items():
    assembled = assemble(source, results[path])
    generator.arrange.accompaniment.audit(original, assembled)
    reuse = generator.arrange.accompaniment.audit(original, {role: value['document'] for role, value in results[path].items()}, role_documents=True)
    for part, document in assembled.items():
      document['metadata'] = copy.deepcopy(original['metadata'])
      document['metadata']['song_id'] = generator.score.slugify(path.stem)
      generator.score.write_json(output_dir / (generator.score.slugify(path.stem) + '_' + part + '.json'), document)
    analysis_dir = output_dir.parent / 'analysis'
    analysis_dir.mkdir(parents=True, exist_ok=True)
    (analysis_dir / (generator.score.slugify(path.stem) + '.txt')).write_text('\n'.join(results[path][r]['report'] for r in ROLES))
    # Not a score file: diagnostics are deliberately outside the songs schema.
    stable = {role: {k: v for k, v in value.items() if k not in ('seconds', 'composition_timings', 'cache_hit', 'observed_reads')} for role, value in results[path].items()}
    for role, audit in reuse.items():
      stable[role]['details']['reuse_audit'] = audit
    (analysis_dir / (generator.score.slugify(path.stem) + '.roles')).write_text(json.dumps(stable, sort_keys=True, indent=2))
  return len(sources) * 3