"""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, score, vocabulary
from generator.arrange import form, generate, pickup
from generator.solo import analysis, composer, notation, report, learned
from generator.solo import authored as repertoire

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', '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 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
    start=cursor
    if section.role == 'slash' and 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-model.fraction(source[0]['duration'])
      a,_,old=result[-1];result[-1]=(a,start,old)
    result.append((start, end, part+'_'+role))
    cursor = end
  return tuple(result)


def dependencies(role, resources):
  names = ('harmony', 'lead_ins') if role.endswith('_solo') else ('accompaniment',) 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 += 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', '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):
  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 learned.Model.load(Path(directory)/'models'/'solo.json')


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):
  from generator.input_audit import composition
  started=time.monotonic()
  with 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=Path(data_dir); resources=vocabulary.load(data_dir/'composition')
  settings=settings or {}
  registered=repertoire.lookup(source,role.removesuffix('_solo'),data_dir) if role.endswith('_solo') and not settings.get('predict',False) else None
  algorithm=VERSION+'/'+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=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=''; composition_timings={}
  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
      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]
      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=accompaniment.carried_chords(result['staves'][0]['measures'])[first-1]
      if registered:
        solo=repertoire.realize(registered,analysis.analyze(document),template,part,start,end,initial)
      else:
        solo_model=_solo_model(str(data_dir),inputs['catalogs']['model'])
        solo=composer.compose(analysis.analyze(document),template,part,end,
          initial=initial,learned_model=solo_model,
          expressions_enabled=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=notation.measures(result['staves'][0]['measures'][notation_first:first+count],relative,'solo_'+part)
      notation.replace_region(result,notation_first,output)
      details=report.structure(solo)
      details['allocation']={'start':str(start),'end':str(end),'lead_in_measures':first-notation_first}
      solo_report=report.describe(solo,first,offset,detail=details)
    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,'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':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, 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=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 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])
    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','composition_timings','cache_hit','observed_reads')} 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