from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor
import hashlib
import json
import multiprocessing
from pathlib import Path
import shutil

from PIL import Image, ImageChops

from generator.engraving.fonts import Fonts
from generator.engraving.cache import renderer_fingerprint, renderer_namespace, restore_pages, score_fingerprint, store_pages
from generator.engraving.layout import layout_score
from generator.engraving.raster import rasterize_page
from generator.engraving.types import EngravingError, Style
from generator.model import validate_document

ENGINE_NAME = 'Vamp Python'
ENGINE_VERSION = '1.0.0'
PAGE_WIDTH = 976
PAGE_HEIGHT = 484

def check_page(path: Path) -> None:
  with Image.open(path) as image:
    image.load()
    if image.format != 'PNG' or image.size != (PAGE_WIDTH, PAGE_HEIGHT):
      raise EngravingError(f'{path.name}: expected a {PAGE_WIDTH}x{PAGE_HEIGHT} PNG page')
    white = Image.new('RGBA', image.size, 'white')
    white.alpha_composite(image.convert('RGBA'))
    ink = ImageChops.difference(white.convert('RGB'), Image.new('RGB', image.size, 'white')).convert('L').point(lambda x: 255 if x > 32 else 0)
    bounds = ink.getbbox()
    if not bounds:
      raise EngravingError(f'{path.name}: engraver produced a blank page')
    if bounds[0] < 4 or bounds[1] < 4 or bounds[2] > PAGE_WIDTH - 4 or bounds[3] > PAGE_HEIGHT - 4:
      raise EngravingError(f'{path.name}: notation reaches the page edge; refusing potentially clipped output')

def engrave_scores(paths: list[Path], root: Path, fonts: Path, cache_dir: Path | None, progress: bool, worker: int, renderer: str) -> tuple[int, int, int, int]:
  family = Fonts(fonts)
  style = Style()
  fingerprint = family.fingerprint(style)
  reused, missing, rejected = 0, 0, 0
  for index, path in enumerate(paths, 1):
    document = json.loads(path.read_text(encoding='utf-8'))
    try:
      validate_document(document)
      destination = path.with_suffix('.pages')
      if destination.exists():
        shutil.rmtree(destination)
      entry = cache_dir / score_fingerprint(document, renderer) if cache_dir else None
      present = entry is not None and (entry.exists() or entry.is_symlink())
      pages = restore_pages(entry, destination) if entry else None
      if pages is None:
        if entry:
          rejected += int(present)
          missing += int(not present)
        if destination.exists():
          shutil.rmtree(destination)
        destination.mkdir()
        layout = layout_score(document, style, family)
        pages = []
        for page in layout.pages:
          name = f'{page.number:04}.png'
          target = destination / name
          rasterize_page(page, family).save(target, format='PNG')
          check_page(target)
          pages.append({'path': name, 'sha256': hashlib.sha256(target.read_bytes()).hexdigest()})
        if entry:
          store_pages(entry, destination, pages)
      else:
        reused += 1
      entries = [{**page, 'path': destination.name + '/' + page['path']} for page in pages]
      music = {k: v for k, v in document.items() if k not in ('schema_version', 'playback')}
      bundle = {'schema_version': 3, 'music': music, 'presentation': {'engine': ENGINE_NAME, 'engine_version': ENGINE_VERSION, 'style_sha256': fingerprint, 'width': PAGE_WIDTH, 'height': 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')
      if progress and (index % 50 == 0 or index == len(paths)):
        print(f'Engraving worker {worker}: {index}/{len(paths)} scores, {reused} reused, {missing} missing, {rejected} rejected', flush=True)
    except ValueError as error:
      raise EngravingError(f'{path.relative_to(root)}: {error}') from error
  return len(paths), reused, missing, rejected

def engrave_tree(root: Path, fonts: Path, *, cache_dir: Path | None = None, jobs: int = 1, progress: bool = False) -> None:
  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')
  paths = sorted(root.rglob('*.json'))
  if not paths:
    return
  renderer = renderer_fingerprint(Fonts(fonts), Style()) if cache_dir else ''
  if cache_dir:
    cache_dir = renderer_namespace(cache_dir.resolve(), 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)
  jobs = min(jobs, len(paths))
  if jobs == 1:
    results = [engrave_scores(paths, root, fonts, cache_dir, progress, 1, renderer)]
  else:
    with ProcessPoolExecutor(max_workers=jobs, mp_context=multiprocessing.get_context('spawn')) as pool:
      futures = [pool.submit(engrave_scores, paths[index::jobs], root, fonts, cache_dir, progress, index + 1, renderer) for index in range(jobs)]
      results = [future.result() for future in futures]
  if progress:
    total, reused, missing, rejected = (sum(row[index] for row in results) for index in range(4))
    print(f'Engraving complete: {total} scores, {reused} reused, {missing} missing, {rejected} rejected', flush=True)