"""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
from pathlib import Path
import tempfile
import time
from fractions import Fraction as F

from generator import accompaniment, comping, fretboard, model, musicxml, parts, rhythm, score, vocabulary
from generator.arrange import form, generate, pickup
from generator.solo import analysis, composer, notation, report, sequences
from generator.scales import library as scales

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 = score.source_staff(prepared)
    for i, m in enumerate(staff['measures']):
      m['id'], m['number'] = f'source_{i}', str(i+1)
    staff['tuplets'] = 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 = 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', 'meter', '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 windows(document, part):
  source = score.source_staff(document)['measures']
  sections = form.full_arrangement_sections(source, part)
  result, cursor = [], F()
  for section in sections:
    duration = sum((model.fraction(m['duration']) for m in source[section.source_measure_index:section.source_measure_index+section.measure_count]), F())
    role = {'slash': 'solo', 'vocal': 'melody'}.get(section.role, section.role)
    if section.role in ('vocal', 'melody') and pickup.enters(section, source) and result:
      attacks = [model.fraction(n['onset']) for n in score.measure_notes(source[0]) if n['pitches']]
      if attacks:
        start = cursor-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
    result.append((cursor, end, part+'_'+role))
    if part == 'electric' and section.role == 'slash' and score.body_start_measure_index(source):
      attacks = [model.fraction(n['onset']) for n in score.measure_notes(source[0]) if n['pitches']]
      if attacks:
        cutoff = end-model.fraction(source[0]['duration'])+min(attacks)
        result[-1] = (cursor, cutoff, part+'_solo')
        result.append((cutoff, end, part+'_rest'))
    cursor = end
  return tuple(result)


def dependencies(role, resources):
  names = ('grammar', 'harmony', 'lead_ins', 'rhythms', 'melodic_vocabulary') if role.endswith('_solo') else ('rhythms',) if role.endswith('_accompaniment') else ()
  return {'roles': {'role': resources['roles']['roles'][role], 'difficulty': resources['roles']['difficulty']},
    **{name: resources[name] for name in names}}


def _code_key(role):
  root = 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 += [p for p in (root/'solo').glob('*.py') if 'audition' not in p.name and p.name != 'references.py']
    paths += list((root/'scales').glob('*.py')) + [root/'rhythm.py']
  elif role.endswith('_accompaniment'):
    paths += [root/n for n in ('comping.py', 'accompaniment.py', 'guitar.py')]
  elif role == 'electric_melody':
    paths += [root/'guitar.py', root/'melody.py']
  return 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):
  names = ('scales', 'sequences', 'rhythms') if role.endswith('_solo') else ()
  return {name: vocabulary.fingerprint([(p.relative_to(directory/name).as_posix(), hashlib.sha256(p.read_bytes()).hexdigest()) for p in sorted((directory/name).rglob('*.xml'))]) for name in names}

@functools.lru_cache(maxsize=4)
def _solo_libraries(directory, fingerprints):
  return scales.load_library(directory/'scales'),sequences.load_library(directory/'sequences'),rhythm.load_library(directory/'rhythms')


def _mask(document, allowed):
  """Apply source-defined role exclusions, without observing another part."""
  absolute = F()
  for measure in document['staves'][0]['measures']:
    length = model.fraction(measure['duration'])
    for voice in measure['voices']:
      selected = []
      for n in voice['notes']:
        start = absolute+model.fraction(n['onset'])
        end = start+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'] = model.rational(a-absolute)
            for pitch in note['pitches']: pitch.pop('tie_from',None)
          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']=pickup.retain_links(staff['measures'],staff.get('tuplets',[]))


def _fill(notes, length, prefix, voice='voice_arranged'):
  notes=sorted(notes,key=lambda n:model.fraction(n['onset']))
  result=[]; cursor=F()
  def rests(end):
    nonlocal cursor
    for duration,written in accompaniment.written_durations(end-cursor):
      result.append(model.make_event(prefix+'_rest_'+str(len(result)),cursor,duration,'rest',written=written)|{'voice':voice})
      cursor+=duration
  for n in notes:
    start=model.fraction(n['onset'])
    if start<cursor: raise ValueError(prefix+': overlapping independently generated roles')
    rests(start); n['voice']=voice;result.append(n)
    cursor=start+model.fraction(n['duration'])
  rests(length)
  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=Path(data_dir); resources=vocabulary.load(data_dir/'composition')
  settings=settings or {}
  from generator.solo.policy import DEFAULT
  algorithm=VERSION+'/'+DEFAULT.composition_version if role.endswith('_solo') else VERSION
  inputs={'version':algorithm,'source':source.musical_key,'role':role,'settings':settings,
    'vocabulary':dependencies(role,resources),'catalogs':_resource_key(data_dir,role),'code':_code_key(role)}
  # Source notation/lyrics also invalidate cached score representation, never the musical seed.
  key=vocabulary.fingerprint({'inputs':inputs,'notation':source.document_json})
  path=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']==vocabulary.fingerprint(value):
        model.validate_document(value['document']); 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':vocabulary.fingerprint(inputs['vocabulary']),'settings':settings,'algorithm':algorithm,'code':inputs['code'],'cache_key':key}
  seed=vocabulary.fingerprint({'source':source.musical_key,'role':role,'settings':settings,'data':inputs['vocabulary'],'version':algorithm})
  details={}; solo_report=''
  with vocabulary.using(resources,role,settings|{'choice_seed':seed}):
    if role.endswith('_solo'):
      result=generate._arrange_document(document,part,{},owned={'slash'})
      from generator.solo.arrangement import regions,annotate
      si=3 if part=='acoustic' else 4
      first,count,offset,length=regions(document)[si]
      # The template/harmony comes directly from source form, not another generated part.
      authored=score.source_staff(document)['measures']
      section=form.full_arrangement_sections(authored,part)[si]
      effective=generate.section_source(authored,section,form.full_arrangement_sections(authored,part)[si+1])
      template=copy.deepcopy(effective[section.source_measure_index:section.source_measure_index+count])
      initial=accompaniment.carried_chords(effective)[section.source_measure_index]
      end=next(b for a,b,r in windows(document,part) if r==role)-offset
      scale_library,sequence_library,rhythm_library=_solo_libraries(data_dir,tuple(sorted(inputs['catalogs'].items())))
      solo=composer.compose(analysis.analyze(document),template,part,end,
        initial=initial,scales=scale_library,sequences=sequence_library,rhythms=rhythm_library,
        enclosures_enabled=vocabulary.role().get('expressions',True))
      output=notation.measures(result['staves'][0]['measures'][first:first+count],solo.notes(),'solo_'+part)
      notation.replace_region(result,first,output);annotate(result,solo,offset)
      details=report.structure(solo)
      solo_report=report.describe(solo,scale_library,first,offset)
    elif role.endswith('_accompaniment'):
      canonical=comping.Canonical(document,part,seed)
      result=generate._arrange_document(document,part,{},owned={'accompaniment'},accompaniment=canonical)
      details={'canonical_comping':canonical.report}
    else:
      result=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)
  staff=result['staves'][0]
  for measure in staff['measures']:
    measure['id']=role+'_'+measure['id']
  staff['tuplets']=score.reidentify(staff['measures'],staff.get('tuplets',[]))
  if part!='vocal':result['staves'][0]['tablature']=True
  model.validate_document(result);fretboard.validate_document(result)
  value={'role':role,'document':result,'trace':trace,'details':details,'report':solo_report,'seconds':time.monotonic()-started,'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':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 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.
      sounding=[copy.deepcopy(n) for d in selected for n in score.measure_notes(d['staves'][0]['measures'][i]) if n['pitches'] or n.get('tuplet_ids')]
      m['voices']=[{'id':'voice_arranged','notes':_fill(sounding,model.fraction(m['duration']),m['id'])}]
      # 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']=pickup.retain_links(staff['measures'],staff['tuplets'])
    for span in staff['tuplets']:span['voice']='voice_arranged'
    model.validate_document(document);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):
  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=Path(source_dir),Path(output_dir)
  sources={}
  for path in parts.source_paths(source_dir):
    document=parts.prepare_document(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 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])
    for part,document in assembled.items():
      document['metadata']=copy.deepcopy(original['metadata'])
      score.write_json(output_dir/(score.slugify(path.stem)+'_'+part+'.json'),document)
    analysis_dir=output_dir.parent/'analysis';analysis_dir.mkdir(parents=True,exist_ok=True)
    (analysis_dir/(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','cache_hit')} for role,value in results[path].items()}
    (analysis_dir/(score.slugify(path.stem)+'.roles')).write_text(json.dumps(stable,sort_keys=True,indent=2))
  return len(sources)*3