from __future__ import annotations

import concurrent.futures
import json
import multiprocessing
import pathlib
import time

import generator.engraving.cache
import generator.engraving.fonts
import generator.engraving.types
import generator.engraving.validation
import generator.engraving.worker
import generator.validation

def engrave_tree(root: pathlib.Path, fonts: pathlib.Path, *, cache_dir: pathlib.Path | None = None, jobs: int = 1, progress: bool = False) -> set[str]:
  if type(jobs) is not int or not 1 <= jobs <= 16:
    raise ValueError('engraving jobs must be between 1 and 16')
  if cache_dir and (cache_dir.resolve().is_relative_to(root.resolve()) or root.resolve().is_relative_to(cache_dir.resolve())):
    raise ValueError('engraving cache and generated assets must not overlap')
  started = time.monotonic()
  paths = sorted(root.rglob('*.json'))
  validated_pages = set()
  if not paths:
    return validated_pages
  family, style = generator.engraving.fonts.Fonts(fonts), generator.engraving.types.Style()
  fingerprint = family.fingerprint(style)
  renderer = generator.engraving.cache.renderer_fingerprint(family, style)
  if cache_dir:
    cache_dir = generator.engraving.cache.renderer_namespace(cache_dir, renderer)
    if progress:
      state = 'warm namespace' if cache_dir.is_dir() else 'cold namespace'
      print(f'Engraving cache enabled: {cache_dir} ({state}; renderer {renderer})', flush=True)
  elif progress:
    print('Engraving cache disabled', flush=True)
  totals = generator.engraving.worker.GroupResult()
  missing, rejected, completed = 0, 0, 0
  groups = {}
  for path in paths:
    try:
      document = json.loads(path.read_text(encoding='utf-8'))
      generator.validation.validate_document(document)
      score_key = generator.engraving.cache.score_fingerprint(document, renderer)
      destination = path.with_suffix('.pages')
      generator.engraving.worker.clear_pages(destination)
      entry = cache_dir / 'fullscore' / score_key if cache_dir else None
      present = entry is not None and (entry.exists() or entry.is_symlink())
      pages = generator.engraving.cache.restore_pages(entry, destination, validated_pages) if entry else None
      if pages is not None:
        try:
          generator.engraving.worker.validate_cached_views(document, pages)
        except (ValueError, TypeError, KeyError):
          pages = None
      if pages is None:
        rejected += int(present)
        missing += int(entry is not None and not present)
        notation = generator.engraving.cache.notation_fingerprint(document, renderer)
        groups.setdefault(notation, []).append((path, score_key))
      else:
        generator.engraving.worker.write_bundle(path, document, pages, fingerprint)
        totals.fullscore += 1
        completed += 1
    except ValueError as error:
      raise generator.engraving.types.EngravingError(f'{path.relative_to(root)}: {error}') from error
  prepare_time = time.monotonic() - started
  if progress:
    print(f'Engraving grouped: {len(paths)} scores, {len(groups)} notation bodies, {totals.fullscore} fullscore hits', flush=True)
  validation_time = 0.0

  def finish(result):
    nonlocal completed, validation_time
    started = time.monotonic()
    for path, score_key, pages in result.scores:
      destination = path.with_suffix('.pages')
      for page in pages:
        generator.engraving.validation.check_page(destination / page['path'], validated_pages, expected_sha256=page['sha256'])
    validation_time += time.monotonic() - started
    if cache_dir:
      started = time.monotonic()
      stored = set()
      for path, score_key, pages in result.scores:
        if score_key not in stored:
          generator.engraving.cache.store_pages(cache_dir / 'fullscore' / score_key, path.with_suffix('.pages'), pages)
          stored.add(score_key)
      result.timings['cache'] += time.monotonic() - started
    for field in ('fullscore', 'bodyhits', 'bodies', 'renderedpages', 'rejected'):
      setattr(totals, field, getattr(totals, field) + getattr(result, field))
    for phase, elapsed in result.timings.items():
      totals.timings[phase] += elapsed
    completed += len(result.scores)
    if progress:
      print(f'Engraving: {completed}/{len(paths)} scores, fullscore {totals.fullscore}, bodyhits {totals.bodyhits}, bodies {totals.bodies}, renderedpages {totals.renderedpages}', flush=True)

  jobs = min(jobs, len(groups))
  if jobs <= 1:
    for notation, group in groups.items():
      finish(generator.engraving.worker.engrave_group(notation, group, root, fonts, cache_dir, fingerprint, progress))
  else:
    with concurrent.futures.ProcessPoolExecutor(max_workers=jobs, mp_context=multiprocessing.get_context('spawn')) as pool:
      futures = [pool.submit(generator.engraving.worker.engrave_group, notation, group, root, fonts, cache_dir, fingerprint, progress) for notation, group in groups.items()]
      for future in concurrent.futures.as_completed(futures):
        finish(future.result())
  if progress:
    print(f'Engraving complete: {len(paths)} scores, {totals.fullscore} reused, {missing} missing, {rejected} rejected; fullscore {totals.fullscore}, bodyhits {totals.bodyhits}, bodies {totals.bodies}, renderedpages {totals.renderedpages}, rejected bodies {totals.rejected}', flush=True)
    phases = {'prepare': prepare_time, **totals.timings, 'validation': validation_time}
    print('Engraving timings: ' + ', '.join(f'{phase} {elapsed:.3f}s' for phase, elapsed in phases.items()) + f'; wall {time.monotonic() - started:.3f}s (worker phases summed)', flush=True)
  return validated_pages