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, 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) -> tuple[int, int]:
  family = Fonts(fonts)
  style = Style()
  fingerprint = family.fingerprint(style)
  renderer = renderer_fingerprint(family, style) if cache_dir else ''
  reused = 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
      pages = restore_pages(entry, destination) if entry else None
      if pages is None:
        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', flush=True)
    except ValueError as error:
      raise EngravingError(f'{path.relative_to(root)}: {error}') from error
  return len(paths), reused

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
  jobs = min(jobs, len(paths))
  if jobs == 1:
    results = [engrave_scores(paths, root, fonts, cache_dir, progress, 1)]
  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) for index in range(jobs)]
      results = [future.result() for future in futures]
  if progress:
    print(f'Engraving complete: {sum(total for total, _ in results)} scores, {sum(reused for _, reused in results)} reused', flush=True)