"""File discovery, process scheduling and publication of independent roles."""
from __future__ import annotations
import concurrent.futures
import copy
import json
import multiprocessing
import pathlib

import generator.arrange.accompaniment
import generator.arrange.assembly
import generator.composition
import generator.composition_types
import generator.musicxml
import generator.parts
import generator.score
import generator.score_io
import generator.source

def _job(request):
  return generator.composition.generate_role(request.source, request.role, request.directory, json.loads(request.settings_json), request.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, generator.source.SourceAnalysis.read(document))
  requests = [(path, role, generator.composition_types.RoleRequest(s, role, source_dir.parent, cache=cache_dir)) for path, (_, s) in sources.items() for role in generator.composition_types.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 = generator.arrange.assembly.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_io.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 generator.composition_types.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