"""Build-only PetalumaScript sprites and content-addressed Pico title images."""

from __future__ import annotations

import argparse
import hashlib
import json
import pathlib
import struct

import PIL.Image
import PIL.ImageDraw
import PIL.ImageFont
import PIL.features
import fontTools.ttLib

TITLE_WIDTH = 296
TITLE_HEIGHT = 64
TITLE_HEADER = struct.pack('<4sHH', b'VPTI', TITLE_WIDTH, TITLE_HEIGHT)
TITLE_BYTES = len(TITLE_HEADER) + TITLE_WIDTH * TITLE_HEIGHT // 2
MAX_TITLES = 8192
LABELS = ('VAMP', 'CONNECTING', 'OFFLINE', 'BUSY', 'READY', 'LOADING', 'WIFI', 'ERROR', 'PAGE', 'Acoustic', 'Electric', 'Vocal', 'UNAVAILABLE',
  'Main Menu', 'Performance', 'Settings', 'Backing', 'Metronome', 'Solos', 'Tempo', 'Back', 'On', 'Off', 'BPM', 'Manual', 'Preparing', 'Count-in', 'Playing', 'Sustain', 'Hold to exit')
SPRITES = LABELS + tuple('0123456789/')
DIGIT_WIDTH = 22
DEFAULT_FONT_ROOT = pathlib.Path(__file__).resolve().parents[1] / 'package/assets/fonts/petaluma'

class Renderer:
  def __init__(self, font_path: pathlib.Path):
    if not PIL.features.check_feature('raqm'):
      raise RuntimeError('Pico artwork requires Pillow with Raqm support')
    self.font_path = font_path
    with fontTools.ttLib.TTFont(font_path) as font:
      self.characters = set(font.getBestCmap())
    self.fonts = {}
  def font(self, size: int) -> PIL.ImageFont.FreeTypeFont:
    if size not in self.fonts:
      self.fonts[size] = PIL.ImageFont.truetype(str(self.font_path), size, layout_engine=PIL.ImageFont.Layout.RAQM)
    return self.fonts[size]
  def validate(self, text: str) -> None:
    if not isinstance(text, str) or not text.strip() or len(text) > 200:
      raise ValueError('Pico title and part must contain 1 to 200 characters')
    if any(ord(character) < 32 or ord(character) not in self.characters for character in text):
      raise ValueError(f'Unsupported PetalumaScript text: {text!r}')
  def line(self, text: str, width: int, height: int, size: int) -> PIL.Image.Image:
    self.validate(text)
    if width < 5 or height < 3:
      raise ValueError('Pico text tile is too small')
    font = self.font(size)
    while size > 14:
      box = font.getbbox(text)
      if box[2] - box[0] <= width - 4 and box[3] - box[1] <= height - 2:
        break
      size -= 1
      font = self.font(size)
    if font.getbbox(text)[2] - font.getbbox(text)[0] > width - 4:
      if font.getbbox('...')[2] - font.getbbox('...')[0] > width - 4:
        raise ValueError('Pico text tile cannot fit an ellipsis')
      original = text
      while original and font.getbbox(original + '...')[2] - font.getbbox(original + '...')[0] > width - 4:
        original = original[:-1]
      text = original + '...'
    box = font.getbbox(text)
    if box[3] - box[1] > height - 2:
      raise ValueError('Pico text tile is too short')
    image = PIL.Image.new('L', (width, height))
    PIL.ImageDraw.Draw(image).text(((width - box[2] + box[0]) // 2 - box[0],
      (height - box[3] + box[1]) // 2 - box[1]), text, font=font, fill=255)
    return image
  def digit(self, text: str) -> PIL.Image.Image:
    font = self.font(24)
    box = font.getbbox(text)
    if text not in '0123456789/' or len(text) != 1 or box[2] - box[0] > DIGIT_WIDTH - 4 or box[1] + 2 < 0 or box[3] + 2 > 28:
      raise ValueError('Invalid Pico counter tile')
    image = PIL.Image.new('L', (DIGIT_WIDTH, 28))
    PIL.ImageDraw.Draw(image).text(((DIGIT_WIDTH - box[2] + box[0]) // 2 - box[0], 2), text, font=font, fill=255)
    return image
  def title(self, title: str) -> bytes:
    self.validate(title)
    for size in range(32, 15, -1):
      lines = ['']
      for word in title.split():
        candidate = (lines[-1] + ' ' + word).strip()
        box = self.font(size).getbbox(candidate)
        if lines[-1] and box[2] - box[0] > TITLE_WIDTH - 4:
          lines.append(word)
        else:
          lines[-1] = candidate
      if len(lines) <= 2:
        break
    if len(lines) > 2:
      lines = [lines[0], ' '.join(lines[1:])]
    image = PIL.Image.new('L', (TITLE_WIDTH, TITLE_HEIGHT))
    for index, line in enumerate(lines):
      image.paste(self.line(line, TITLE_WIDTH, 32, size), (0, (TITLE_HEIGHT - len(lines) * 32) // 2 + index * 32))
    return TITLE_HEADER + pack_coverage(image)

def pack_coverage(image: PIL.Image.Image) -> bytes:
  pixels = image.tobytes()
  if len(pixels) % 2:
    pixels += b'\0'
  levels = [(pixel + 8) // 17 for pixel in pixels]
  return bytes((levels[index] << 4) | levels[index + 1] for index in range(0, len(levels), 2))

def generate_static(output: pathlib.Path, font_root: pathlib.Path = DEFAULT_FONT_ROOT) -> None:
  renderer = Renderer(font_root / 'PetalumaScript.otf')
  directory = bytearray(struct.pack('<4sHH', b'VPSA', len(SPRITES), 0))
  payload = bytearray()
  for text in SPRITES:
    image = renderer.line(text, 176, 28, 24) if text in LABELS else renderer.digit(text)
    if text in LABELS:
      box = image.getbbox()
      image = image.crop((max(0, box[0] - 1), 0, min(image.width, box[2] + 1), image.height))
    directory.extend(struct.pack('<HHI', image.width, image.height, 8 + 8 * len(SPRITES) + len(payload)))
    payload.extend(pack_coverage(image))
  output.parent.mkdir(parents=True, exist_ok=True)
  output.write_bytes(directory + payload)

def generate_titles(root: pathlib.Path, font_root: pathlib.Path = DEFAULT_FONT_ROOT) -> None:
  """Call after score generation, before catalog validation and atomic promotion.

  Adds presentation.title_image to score bundles. The identifier is SHA-256 of
  the complete fixed-size .mask file, shared across parts with the same title.
  """
  renderer = Renderer(font_root / 'PetalumaScript.otf')
  destination = root / 'pico' / 'titles'
  destination.mkdir(parents=True, exist_ok=True)
  titles = {}
  for source in sorted((root / 'songs').rglob('*.json')):
    document = json.loads(source.read_text(encoding='utf-8'))
    title = document['music']['metadata']['title']
    if title not in titles:
      if len(titles) >= MAX_TITLES:
        raise ValueError('Too many Pico titles')
      image = renderer.title(title)
      titles[title] = hashlib.sha256(image).hexdigest()
      (destination / f'{titles[title]}.mask').write_bytes(image)
    document['presentation']['title_image'] = titles[title]
    source.write_text(json.dumps(document, ensure_ascii=False, separators=(',', ':')) + '\n', encoding='utf-8')

def validate_title(root: pathlib.Path, identifier: str) -> None:
  if not isinstance(identifier, str) or len(identifier) != 64 or any(character not in '0123456789abcdef' for character in identifier):
    raise ValueError('Invalid Pico title identifier; regenerate assets')
  try:
    asset_root = root.resolve(strict=True)
    directory = (root / 'pico' / 'titles').resolve(strict=True)
    source = (directory / f'{identifier}.mask').resolve(strict=True)
    if not directory.is_relative_to(asset_root) or source.parent != directory:
      raise ValueError('Pico title escapes the generated title directory')
    if source.stat().st_size != TITLE_BYTES:
      raise ValueError('Invalid Pico title size')
    with source.open('rb') as stream:
      image = stream.read(TITLE_BYTES + 1)
    if len(image) != TITLE_BYTES or image[:len(TITLE_HEADER)] != TITLE_HEADER or hashlib.sha256(image).hexdigest() != identifier:
      raise ValueError('Invalid Pico title header or digest')
  except OSError as error:
    raise ValueError(f'Pico title unavailable: {error}') from error

def main() -> None:
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument('--font-root', type=pathlib.Path, default=DEFAULT_FONT_ROOT)
  group = parser.add_mutually_exclusive_group(required=True)
  group.add_argument('--static-output', type=pathlib.Path)
  group.add_argument('--assets', type=pathlib.Path)
  arguments = parser.parse_args()
  if arguments.static_output:
    generate_static(arguments.static_output, arguments.font_root)
  else:
    generate_titles(arguments.assets, arguments.font_root)

if __name__ == '__main__':
  main()