"""Independent, reproducible role generation from source and policy inputs.

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 copy
import dataclasses
import functools
import hashlib
import json
import pathlib
import time

import generator.accompaniment
import generator.acoustic
import generator.arrange.assembly
import generator.arrange.form
import generator.arrange.generate
import generator.comping
import generator.comping_rhythm
import generator.composition_types
import generator.fretboard
import generator.input_audit
import generator.role_cache
import generator.score
import generator.solo.arrangement
import generator.solo.authored
import generator.solo.composer
import generator.solo.learned
import generator.solo.notation
import generator.solo.report
import generator.source
import generator.validation
import generator.vocabulary

VERSION = 'independent-roles-1'

def dependencies(role, resources, measures=None):
  names = ['harmony', 'lead_ins'] if role.endswith('_solo') else []
  accompaniment = {}
  if role.endswith('_accompaniment'):
    policy = generator.vocabulary.accompaniment(role.split('_')[0], resources)
    accompaniment = {'accompaniment': generator.comping_rhythm.relevant_policy(policy, measures) if measures is not None else policy}
  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,
    **{name: resources[name] for name in names}}

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 generate_role(source, role, data_dir, settings=None, cache_dir=None) -> generator.composition_types.RoleResult:
  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 generator.composition_types.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, generator.score.source_staff(source.document())['measures']), 'catalogs': ({'authored': registered['sha256'], 'id': registered['id'], 'pitch_convention': registered['pitch_convention']} if registered else _resource_key(data_dir, role)), 'code': generator.role_cache.code_key()}
  # 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
  cached = generator.role_cache.restore(path, key)
  if cached is not None:
    return cached
  started = time.monotonic()
  document = source.document()
  part = role.split('_')[0]
  plan = generator.arrange.form.FormPlan.read(document, part)
  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)
    for field in ('provenance', 'license', 'research'):
      choice_data['accompaniment'].pop(field, 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 = {}
  composition = generator.vocabulary.CompositionContext.create(resources, role, settings | {'choice_seed': seed})
  if role.endswith('_solo'):
    result = generator.arrange.generate._arrange_document(document, part, plan, owned={'slash'}, composition=composition)
    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 = plan.sections[si]
    effective = generator.arrange.generate.section_source(authored, section, plan.sections[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 plan.windows 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.source.analyze(document), template, part, start, end, initial, composition=composition)
    else:
      solo_model = _solo_model(str(data_dir), inputs['catalogs']['model'])
      solo = generator.solo.composer.compose(generator.source.analyze(document), template, part, end, initial=initial, learned_model=solo_model, expressions_enabled=composition.role().get('expressions', True), start=start, composition=composition)
    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, composition=composition)
    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, composition=composition)
  elif role.endswith('_accompaniment'):
    canonical = generator.comping.Canonical(document, part, seed, composition=composition, plan=plan)
    result = generator.arrange.generate._arrange_document(document, part, plan, owned={'accompaniment'}, accompaniment=canonical, composition=composition)
    details = {'canonical_comping': canonical.report, 'guitar_placement': canonical.placement_report, 'accompaniment_rhythm': canonical.rhythm_report, 'groove_cycles': canonical.groove_report, 'rhythm_source': choice_source}
  else:
    placement_report = {}
    result = generator.arrange.generate._arrange_document(document, part, plan, owned={'vocal' if part == 'vocal' else 'melody'}, composition=composition, placement_report=placement_report)
    if placement_report:
      details['guitar_placement'] = placement_report
  allowed = tuple((a, b) for a, b, r in plan.windows if r == role)
  generator.arrange.assembly.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.validation.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))
  generator.role_cache.store(path, key, value)
  return value