from __future__ import annotations

import collections
import concurrent.futures
import dataclasses
import hashlib
import json
import multiprocessing
import pathlib
import shutil
import tempfile
import time

import PIL.Image

import generator.engraving.audit
import generator.engraving.cache
import generator.engraving.fonts
import generator.engraving.layout
import generator.engraving.raster
import generator.engraving.types
import generator.engraving.validation
import generator.model

@dataclasses.dataclass
class GroupResult:
  scores: list = dataclasses.field(default_factory=list)
  fullscore: int = 0
  bodyhits: int = 0
  bodies: int = 0
  renderedpages: int = 0
  rejected: int = 0
  timings: dict = dataclasses.field(default_factory=lambda: dict.fromkeys(('layout', 'pagination', 'raster', 'cache'), 0.0))

class BodyBands:
  def __init__(self, root, persistent, result, limit=16 * 1024 * 1024):
    self.root, self.persistent, self.result, self.limit = root, persistent, result, limit
    self.images = collections.OrderedDict()
    self.size = 0
    self.known = {}
    self.root.mkdir(parents=True, exist_ok=True)
  def get(self, notation, page, band, fonts):
    key = generator.engraving.cache.band_fingerprint(notation, page, band)
    if key in self.images:
      self.images.move_to_end(key)
      self.result.bodyhits += 1
      return self.images[key]
    started = time.monotonic()
    entry = self.root / (key + '.rgb')
    image = None
    if key in self.known:
      if entry.is_symlink() or entry.stat().st_size != page.width * (band.bottom - band.top) * 3:
        raise generator.engraving.types.EngravingError('temporary notation band was corrupted')
      data = entry.read_bytes()
      if hashlib.sha256(data).hexdigest() != self.known[key]:
        raise generator.engraving.types.EngravingError('temporary notation band was corrupted')
      image = PIL.Image.frombytes('RGB', (page.width, band.bottom - band.top), data)
    elif self.persistent:
      cached = self.persistent / 'notation' / key
      image = generator.engraving.cache.restore_band(cached, key, page.width, band.bottom - band.top)
      self.result.rejected += int(image is None and (cached.exists() or cached.is_symlink()))
    self.result.timings['cache'] += time.monotonic() - started
    if image is None:
      started = time.monotonic()
      image = generator.engraving.raster.rasterize_band(page, band.primitives, band.top, band.bottom, fonts)
      self.result.timings['raster'] += time.monotonic() - started
      self.result.bodies += 1
      if self.persistent:
        started = time.monotonic()
        generator.engraving.cache.store_band(self.persistent / 'notation' / key, key, image)
        self.result.timings['cache'] += time.monotonic() - started
    else:
      self.result.bodyhits += 1
    if key not in self.known:
      started = time.monotonic()
      # Spill finalized pixels, so eviction never repeats rasterization or PNG decoding.
      data = image.tobytes()
      entry.write_bytes(data)
      self.result.timings['cache'] += time.monotonic() - started
      self.known[key] = hashlib.sha256(data).hexdigest()
    size = image.width * image.height * 3
    while self.images and self.size + size > self.limit:
      _, old = self.images.popitem(last=False)
      self.size -= old.width * old.height * 3
      old.close()
    self.images[key] = image
    self.size += size
    return image
  def close(self):
    for image in self.images.values():
      image.close()
    self.images.clear()
    self.size = 0

def clear_pages(destination):
  if destination.is_symlink() or destination.is_file():
    destination.unlink()
  elif destination.exists():
    shutil.rmtree(destination)

def write_bundle(path, document, pages, fingerprint):
  entries = [{**page, 'path': path.with_suffix('.pages').name + '/' + page['path']} for page in pages]
  music = {key: value for key, value in document.items() if key != 'playback'}
  bundle = {'music': music, 'presentation': {'engine': generator.engraving.validation.ENGINE_NAME, 'style_sha256': fingerprint,
    'width': generator.engraving.validation.PAGE_WIDTH, 'height': generator.engraving.validation.PAGE_HEIGHT, 'pages': entries}}
  if 'playback' in document:
    bundle['playback'] = document['playback']
  path.write_text(json.dumps(bundle, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')

def engrave_group(notation, paths, root, font_path, cache_dir, fingerprint, progress=False):
  family = generator.engraving.fonts.Fonts(font_path)
  style = generator.engraving.types.Style()
  systems = None
  scores = {}
  audited = set()
  result = GroupResult()
  with tempfile.TemporaryDirectory(prefix='vamp-notation-') as temporary:
    bodies = BodyBands(pathlib.Path(temporary), cache_dir, result)
    try:
      for path, score_key in paths:
        try:
          document = json.loads(path.read_text(encoding='utf-8'))
          style = generator.engraving.types.score_style(document, generator.engraving.types.Style())
          destination = path.with_suffix('.pages')
          clear_pages(destination)
          destination.mkdir()
          if score_key in scores:
            source, pages = scores[score_key]
            for page in pages:
              shutil.copyfile(source / page['path'], destination / page['path'])
            result.fullscore += 1
          else:
            if systems is None:
              started = time.monotonic()
              systems = generator.engraving.layout.layout_systems(document, style, family)
              result.timings['layout'] += time.monotonic() - started
            started = time.monotonic()
            pages_layout = generator.engraving.layout.paginate(systems, document, family, style)
            placement = tuple((page.width, page.height, tuple((system.start, system.end, system.box) for system in page.systems)) for page in pages_layout)
            if placement not in audited:
              layout = generator.engraving.types.LayoutDocument(pages_layout, style, fingerprint)
              generator.engraving.audit.audit_layout(document, layout, family)
              audited.add(placement)
            else:
              for page in pages_layout:
                decoration = tuple(p for p in page.primitives if p.role in ('title', 'subtitle', 'page-number'))
                generator.engraving.audit.audit_primitives(decoration, page.width, page.height, family)
            result.timings['pagination'] += time.monotonic() - started
            pages = []
            for page in pages_layout:
              bands = generator.engraving.raster.page_bands(page)
              if bands is None:
                started = time.monotonic()
                image = generator.engraving.raster.rasterize_page(page, family)
              else:
                body = bodies.get(notation, page, bands[1], family)
                started = time.monotonic()
                image = generator.engraving.raster.compose_page(page, family, bands, body)
              target = destination / f'{page.number:04}.png'
              with image:
                image.save(target, format='PNG')
              result.timings['raster'] += time.monotonic() - started
              result.renderedpages += 1
              pages.append({'path': target.name, 'sha256': hashlib.sha256(target.read_bytes()).hexdigest()})
            scores[score_key] = (destination, pages)
          write_bundle(path, document, pages, fingerprint)
          result.scores.append((path, score_key, pages))
          if progress and len(result.scores) % 50 == 0:
            print(f'Engraving body {notation[:12]}: {len(result.scores)}/{len(paths)} scores, fullscore {result.fullscore}, bodyhits {result.bodyhits}, bodies {result.bodies}, renderedpages {result.renderedpages}', flush=True)
        except ValueError as error:
          raise generator.engraving.types.EngravingError(f'{path.relative_to(root)}: {error}') from error
    finally:
      bodies.close()
  return result

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 = GroupResult()
  missing, rejected, completed = 0, 0, 0
  groups = {}
  for path in paths:
    try:
      document = json.loads(path.read_text(encoding='utf-8'))
      generator.model.validate_document(document)
      score_key = generator.engraving.cache.score_fingerprint(document, renderer)
      destination = path.with_suffix('.pages')
      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 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:
        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(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(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