from __future__ import annotations

import collections
import dataclasses
import hashlib
import json
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.notation
import generator.engraving.raster
import generator.engraving.types
import generator.engraving.validation
import generator.performance
import generator.validation

@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 = [{k: v for k, v in {**page, 'path': path.with_suffix('.pages').name + '/' + page['path']}.items() if k != 'playback'} for page in pages]
  normal = [entry for entry, page in zip(entries, pages) if not page['playback']]
  playback = [entry for entry, page in zip(entries, pages) if page['playback']]
  music = document
  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': normal, 'playback_pages': playback}}
  validate_cached_views(document, pages)
  bundle['performance'] = generator.performance.timing(document)
  path.write_text(json.dumps(bundle, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')

def validate_cached_views(document, pages):
  generator.performance.validate_views(document, [p for p in pages if not p['playback']], [p for p in pages if p['playback']])

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
  index = 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()
              generator.validation.validate_document(document)
              index = generator.engraving.notation.ScoreIndex(document)
              systems = generator.engraving.layout.layout_systems(document, style, family, index=index)
              result.timings['layout'] += time.monotonic() - started
            started = time.monotonic()
            pages_layout, playback_layout = generator.engraving.layout.page_views(systems, document, family, style, index=index)
            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)
              generator.engraving.audit.audit_playback(pages_layout, playback_layout, style, family)
              audited.add(placement)
            else:
              for page in pages_layout + playback_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 = []
            rendered = {}
            for is_playback, page in [(False, p) for p in pages_layout] + [(True, p) for p in playback_layout]:
              if page.primitives in rendered:
                name, digest = rendered[page.primitives]
                pages.append(dict(path=name, sha256=digest, playback=is_playback, **generator.performance.page_timing(document, page)))
                continue
              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 / (('playback-' if is_playback else '') + f'{page.number:04}.png')
              with image:
                image.save(target, format='PNG')
              result.timings['raster'] += time.monotonic() - started
              result.renderedpages += 1
              digest = hashlib.sha256(target.read_bytes()).hexdigest()
              rendered[page.primitives] = target.name, digest
              pages.append(dict(path=target.name, sha256=digest, playback=is_playback, **generator.performance.page_timing(document, page)))
            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