#!/usr/bin/python2
# -*- coding: utf-8 -*-
# anicursorgen
# Copyright (C) 2015 Руслан Ижбулатов <lrn1986@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
from __future__ import print_function

import sys
import os
import argparse
import shlex
import io
import struct
import math
from PIL import Image
from PIL import ImageFilter

p = struct.pack

program_name = 'anicursorgen'
program_version = 1.0

def main ():
  parser = argparse.ArgumentParser (description='Creates .ani or .cur files from separate images and input metadata.',
                        add_help=False)
  parser.add_argument ('-V', '--version', action='version', version='{}-{}'.format (program_name, program_version),
                       help='Display the version number and exit.')
  parser.add_argument ('-h', '-?', action='help',
                       help='Display the usage message and exit.')
  parser.add_argument ('-p', '--prefix', metavar='dir', default=None,
                       help='Find cursor images in the directory specified by dir. If not specified, the current directory is used.')
  parser.add_argument ('-s', '--add-shadows', action='store_true', dest='add_shadows', default=False,
                       help='Generate shadows for cursors (disabled by default).')
  parser.add_argument ('-n', '--no-shadows', action='store_false', dest='add_shadows', default=False,
                       help='Do not generate shadows for cursors (put after --add-shadows to cancel its effect).')

  shadows = parser.add_argument_group (title='Shadow generation', description='Only relevant when --add-shadows is given')

  shadows.add_argument ('-r', '--right-shift', metavar='%', type=float, default=9.375,
                       help='Shift shadow right by this percentage of the canvas size (default is 9.375).')
  shadows.add_argument ('-d', '--down-shift', metavar='%', type=float, default=3.125,
                       help='Shift shadow down by this percentage of the canvas size (default is 3.125).')
  shadows.add_argument ('-b', '--blur', metavar='%', type=float, default=3.125,
                       help='Blur radius, in percentage of the canvas size (default is 3.125, set to 0 to disable blurring).')
  shadows.add_argument ('-c', '--color', metavar='%', default='0x00000040',
                       help='Shadow color in 0xRRGGBBAA form (default is 0x00000040).')

  parser.add_argument ('input_config', default='-', metavar='input-config [output-file]', nargs='?',
                       help='Input config file (stdin by default).')
  parser.add_argument ('output_file', default='-', metavar='', nargs='?',
                       help='Output cursor file (stdout by default).')

  args = parser.parse_args ()

  try:
    if args.color[0] != '0' or args.color[1] not in ['x', 'X'] or len (args.color) != 10:
      raise ValueError
    args.color = (int (args.color[2:4], 16), int (args.color[4:6], 16), int (args.color[6:8], 16), int (args.color[8:10], 16))
  except:
    print ("Can't parse the color '{}'".format (args.color), file=sys.stderr)
    parser.print_help ()
    return 1

  if args.prefix is None:
    args.prefix = os.getcwd ()

  if args.input_config == '-':
    input_config = sys.stdin
  else:
    input_config = open (args.input_config, 'rb')

  if args.output_file == '-':
    output_file = sys.stdout
  else:
    output_file = open (args.output_file, 'wb')

  result = make_cursor_from (input_config, output_file, args)

  input_config.close ()
  output_file.close ()

  return result

def make_cursor_from (inp, out, args):
  frames = parse_config_from (inp, args.prefix)

  animated = frames_have_animation (frames)

  if animated:
    result = make_ani (frames, out, args)
  else:
    buf = make_cur (frames, args)
    copy_to (out, buf)
    result = 0

  return result

def copy_to (out, buf):
  buf.seek (0, io.SEEK_SET)
  while True:
    b = buf.read (1024)
    if len (b) == 0:
      break
    out.write (b)

def frames_have_animation (frames):
  sizes = set ()

  for frame in frames:
    if frame[4] == 0:
      continue
    if frame[0] in sizes:
      return True
    sizes.add (frame[0])

  return False

def make_cur (frames, args):
  buf = io.BytesIO ()
  buf.write (p ('<HHH', 0, 2, len (frames)))
  frame_offsets = []
  for frame in frames:
    width = frame[0]
    if width > 255:
      width = 0
    height = width
    buf.write (p ('<BBBB HH', width, height, 0, 0, frame[1], frame[2]))
    size_offset_pos = buf.seek (0, io.SEEK_CUR)
    buf.write (p ('<II', 0, 0))
    frame_offsets.append ([size_offset_pos])

  for i, frame in enumerate (frames):
    frame_offset = buf.seek (0, io.SEEK_CUR)
    frame_offsets[i].append (frame_offset)

    frame_png = Image.open (frame[3])

    if args.add_shadows:
      succeeded, shadowed = create_shadow (frame_png, args)
      if succeeded == 0:
        frame_png.close ()
        frame_png = shadowed

    compressed = frame[0] > 48

    if compressed:
      write_png (buf, frame, frame_png)
    else:
      write_cur (buf, frame, frame_png)

    frame_png.close ()

    frame_end = buf.seek (0, io.SEEK_CUR)
    frame_offsets[i].append (frame_end - frame_offset)

  for frame_offset in frame_offsets:
    buf.seek (frame_offset[0])
    buf.write (p ('<II', frame_offset[2], frame_offset[1]))

  return buf

def make_framesets (frames):
  framesets = []
  sizes = set ()

  # This assumes that frames are sorted
  size = 0
  for i, frame in enumerate (frames):
    if size == 0 or frame[0] != size:
      size = frame[0]
      counter = 0

      if size in sizes:
        print ("Frames are not sorted: frame {} has size {}, but we have seen that already".format (i, size), file=sys.stderr)
        return None

      sizes.add (size)

    if counter >= len (framesets):
      framesets.append ([])

    framesets[counter].append (frame)
    counter += 1

  for i in range (1, len (framesets)):
    if len (framesets[i - 1]) != len (framesets[i]):
      print ("Frameset {} has size {}, expected {}".format (i, len (framesets[i]), len (framesets[i - 1])), file=sys.stderr)
      return None

  for frameset in framesets:
    for i in range (1, len (frameset)):
      if frameset[i - 1][4] != frameset[i][4]:
        print ("Frameset {} has duration {} for framesize {}, but {} for framesize {}".format (i, frameset[i][4], frameset[i][0], frameset[i - 1][4], frameset[i - 1][0]), file=sys.stderr)
        return None

  return framesets

def make_ani (frames, out, args):
  framesets = make_framesets (frames)
  if framesets is None:
    return 1

  buf = io.BytesIO ()

  buf.write (b'RIFF')
  riff_len_pos = buf.seek (0, io.SEEK_CUR)
  buf.write (p ('<I', 0))
  riff_len_start = buf.seek (0, io.SEEK_CUR)

  buf.write (b'ACON')
  buf.write (b'anih')
  buf.write (p ('<IIIIIIIIII', 36, 36, len (framesets), len (framesets), 0, 0, 32, 1, framesets[0][0][4], 0x01))

  rates = set ()
  for frameset in framesets:
    rates.add (frameset[0][4])

  if len (rates) != 1:
    buf.write (b'rate')
    buf.write (p ('<I', len (framesets) * 4))
    for frameset in framesets:
      buf.write (p ('<I', frameset[0][4]))

  buf.write (b'LIST')
  list_len_pos = buf.seek (0, io.SEEK_CUR)
  buf.write (p ('<I', 0))
  list_len_start = buf.seek (0, io.SEEK_CUR)

  buf.write (b'fram')

  for frameset in framesets:
    buf.write (b'icon')
    cur = make_cur (frameset, args)
    cur_size = cur.seek (0, io.SEEK_END)
    aligned_cur_size = cur_size
    #if cur_size % 4 != 0:
    #  aligned_cur_size += 4 - cur_size % 2
    buf.write (p ('<i', cur_size))
    copy_to (buf, cur)
    pos = buf.seek (0, io.SEEK_END)
    if pos % 2 != 0:
      buf.write ('\x00' * (2 - (pos % 2)))

  end_at = buf.seek (0, io.SEEK_CUR)
  buf.seek (riff_len_pos, io.SEEK_SET)
  buf.write (p ('<I', end_at - riff_len_start))
  buf.seek (list_len_pos, io.SEEK_SET)
  buf.write (p ('<I', end_at - list_len_start))

  copy_to (out, buf)

  return 0

def write_png (out, frame, frame_png):
  frame_png.save (out, "png", optimize=True)

def write_cur (out, frame, frame_png):
  pixels = frame_png.load ()

  out.write (p ('<I II HH IIIIII', 40, frame[0], frame[0] * 2, 1, 32, 0, 0, 0, 0, 0, 0))

  for y in reversed (range (frame[0])):
    for x in range (frame[0]):
      pixel = pixels[x, y]
      out.write (p ('<BBBB', pixel[2], pixel[1], pixel[0], pixel[3]))

  acc = 0
  acc_pos = 0
  for y in reversed (range (frame[0])):
    wrote = 0
    for x in range (frame[0]):
      if pixels[x, y][3] <= 127:
        acc = acc | (1 << acc_pos)
      acc_pos += 1
      if acc_pos == 8:
        acc_pos = 0
        out.write (chr (acc))
        wrote += 1
    if wrote % 4 != 0:
      out.write (b'\x00' * (4 - wrote % 4))

def parse_config_from (inp, prefix):
  frames = []

  for line in inp.readlines ():
    words = shlex.split (line.rstrip ('\n').rstrip ('\r'))

    if len (words) < 4:
      continue

    try:
      size = int (words[0])
      hotx = int (words[1]) - 1
      hoty = int (words[2]) - 1
      filename = words[3]
      if not os.path.isabs (filename):
        filename = prefix + '/' + filename
    except:
      continue

    if len (words) > 4:
      try:
        duration = int (words[4])
      except:
        continue
    else:
      duration = 0

    frames.append ((size, hotx, hoty, filename, duration))

  return frames

def create_shadow (orig, args):
  blur_px = orig.size[0] / 100.0 * args.blur
  right_px = int (orig.size[0] / 100.0 * args.right_shift)
  down_px = int (orig.size[1] / 100.0 * args.down_shift)

  shadow = Image.new ('RGBA', orig.size, (0, 0, 0, 0))
  shadowize (shadow, orig, args.color)
  shadow.load ()

  if args.blur > 0:
    crop = (int (math.floor (-blur_px)), int (math.floor (-blur_px)), orig.size[0] + int (math.ceil (blur_px)), orig.size[1] + int (math.ceil (blur_px)))
    right_px += int (math.floor (-blur_px))
    down_px += int (math.floor (-blur_px))
    shadow = shadow.crop (crop)
    flt = ImageFilter.GaussianBlur (blur_px)
    shadow = shadow.filter (flt)
  shadow.load ()

  shadowed = Image.new ('RGBA', orig.size, (0, 0, 0, 0))
  shadowed.paste (shadow, (right_px, down_px))
  shadowed.crop ((0, 0, orig.size[0], orig.size[1]))
  shadowed = Image.alpha_composite (shadowed, orig)

  return 0, shadowed

def shadowize (shadow, orig, color):
  o_pxs = orig.load ()
  s_pxs = shadow.load ()
  for y in range (orig.size[1]):
    for x in range (orig.size[0]):
      o_px = o_pxs[x, y]
      if o_px[3] > 0:
        s_pxs[x, y] = (color[0], color[1], color[2], int (color[3] * (o_px[3] / 255.0)))

if __name__ == '__main__':
  sys.exit (main ())