#!/usr/bin/env python3

#   GIMP plug-in for JBIG2-encoded PDF files
#   Copyright (C) 2024  Lee Yingtong Li (RunasSudo)
#
#   Loosely adapted from file-openraster.py - Copyright (C) 2009 by Jon Nordby <jononor@gmail.com>, licensed under the GPLv3
#   In turn based on MyPaint source code by Martin Renold
#
#   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 <https://www.gnu.org/licenses/>.

import gi
gi.require_version('Gimp', '3.0')
from gi.repository import Gimp
from gi.repository import GObject
from gi.repository import GLib
from gi.repository import Gio

import os
import subprocess
import sys
import tempfile

def export_jbig2pdf(procedure, run_mode, image, file, options, metadata, config, data):
	# Check jbig2enc version
	try:
		jbig2version = subprocess.run(['jbig2', '-V'], capture_output=True)
	except FileNotFoundError:
		Gimp.message('jbig2 command (from jbig2enc) not in PATH')
		return Gimp.ValueArray.new_from_values([GObject.Value(Gimp.PDBStatusType, Gimp.PDBStatusType.EXECUTION_ERROR)])
	
	if jbig2version.stderr != b'jbig2enc 0.28\n':
		# We only support 0.28 for now, because from 0.29 pdf.py is renamed to jbig2topdf.py
		Gimp.message('Unsupported version of jbig2enc (expected 0.28, got "{}")'.format(jbig2version.stderr.decode('utf-8').strip()))
		return Gimp.ValueArray.new_from_values([GObject.Value(Gimp.PDBStatusType, Gimp.PDBStatusType.EXECUTION_ERROR)])
	
	# jbig2enc OK!
	
	Gimp.progress_init('Exporting PDF JBIG2 image')
	tempdir = tempfile.mkdtemp('gimp-plugin-file-jbig2pdf')
	
	def unroll_layers(layers):
		for layer in layers:
			if not layer.is_group():
				yield layer
			else:
				for sublayer in unroll_layers(layer.get_children()):
					yield sublayer
	
	pngs_out = []
	
	# Export image layer by layer (bottom layer first)
	for i, layer in enumerate(reversed(list(unroll_layers(image.get_layers())))):
		png_out = os.path.join(tempdir, '{:03d}.png'.format(i))
		
		offsets = layer.get_offsets()
		
		# Save layer to new GimpImage
		tmp_img = Gimp.Image.new(image.get_width(), image.get_height(), image.get_base_type())
		tmp_layer = Gimp.Layer.new_from_drawable(layer, tmp_img)
		tmp_img.insert_layer(tmp_layer, None, 0)
		tmp_layer.merge_filters()  # Apply non-destructive effects
		tmp_layer.flatten()  # Remove alpha channel as this confuses jbig2enc
		tmp_img.crop(layer.get_width(), layer.get_height(), layer.get_offsets().offset_x, layer.get_offsets().offset_y)  # Apply layer origin to image origin
		
		# Export tmp_image as PNG
		pdb_proc = Gimp.get_pdb().lookup_procedure('file-png-export')
		pdb_config = pdb_proc.create_config()
		pdb_config.set_property('run-mode', Gimp.RunMode.NONINTERACTIVE)
		pdb_config.set_property('image', tmp_img)
		pdb_config.set_property('file', Gio.File.new_for_path(png_out))
		pdb_config.set_property('options', None)
		pdb_config.set_property('interlaced', 0)   # As per file-openraster.py
		pdb_config.set_property('compression', 2)  # As per file-openraster.py
		pdb_config.set_property('bkgd', True)      # Write all PNG chunks except oFFs(ets) - GIMP defaults
		pdb_config.set_property('offs', False)
		pdb_config.set_property('phys', True)
		pdb_config.set_property('time', True)
		pdb_config.set_property('save-transparent', True)
		pdb_proc.run(pdb_config)
		
		tmp_img.delete()  # Clean up buffer
		
		if os.path.exists(png_out):
			pngs_out.append(png_out)
		else:
			Gimp.message('Unknown error exporting layer {}'.format(i))
			os.rmdir(tempdir)
			Gimp.progress_end()
			return Gimp.ValueArray.new_from_values([GObject.Value(Gimp.PDBStatusType, Gimp.PDBStatusType.EXECUTION_ERROR)])
	
	# Execute jbig2enc
	try:
		subprocess.run(['jbig2', '-s', '-p', '-v'] + pngs_out, cwd=tempdir, check=True)
	except subprocess.CalledProcessError:
		Gimp.message('Unknown error in jbig2enc')
		os.rmdir(tempdir)
		Gimp.progress_end()
		return Gimp.ValueArray.new_from_values([GObject.Value(Gimp.PDBStatusType, Gimp.PDBStatusType.EXECUTION_ERROR)])
	
	# Execute pdf.py
	# Use .tmpsave extension, so we don't overwrite a valid file if there is an exception
	pdf_tmpout = file.peek_path() + '.tmpsave'
	try:
		with open(pdf_tmpout, 'wb') as pdf_file:
			subprocess.run(['pdf.py', 'output'], cwd=tempdir, check=True, stdout=pdf_file)
			success = True
	except subprocess.CalledProcessError:
		success = False
		Gimp.message('Unknown error in jbig2enc pdf.py')
		os.rmdir(tempdir)
		Gimp.progress_end()
		return Gimp.ValueArray.new_from_values([GObject.Value(Gimp.PDBStatusType, Gimp.PDBStatusType.EXECUTION_ERROR)])
	finally:
		if not success:
			os.unlink(pdf_tmpout)
	
	# Success! - Finalise and clean up
	
	# Remove intermediate PNG files
	for png_out in pngs_out:
		os.unlink(png_out)
	
	# Remove intermediate JBIG2 files
	for i in range(len(pngs_out)):
		os.unlink(os.path.join(tempdir, 'output.{:04d}'.format(i)))
	
	# Remove intermediate JBIG2 symbol file and temp directory
	os.unlink(os.path.join(tempdir, 'output.sym'))
	os.rmdir(tempdir)
	
	# Save over destination file
	os.rename(pdf_tmpout, file.peek_path())
	
	Gimp.progress_end()
	return Gimp.ValueArray.new_from_values([GObject.Value(Gimp.PDBStatusType, Gimp.PDBStatusType.SUCCESS)])

class FileJBIG2PDF(Gimp.PlugIn):
	## GimpPlugIn virtual methods ##
	def do_set_i18n(self, procname):
		return False
	
	def do_query_procedures(self):
		return ['file-jbig2pdf-export']
	
	def do_create_procedure(self, name):
		if name == 'file-jbig2pdf-export':
			procedure = Gimp.ExportProcedure.new(self, name, Gimp.PDBProcType.PLUGIN, False, export_jbig2pdf, None)
			procedure.set_image_types('*')
			procedure.set_documentation(
				'save a PDF file with JBIG2 encoding',
				'save a PDF file with JBIG2 encoding',
				name
			)
			procedure.set_menu_label('Portable Document Format (JBIG2)')
			procedure.set_extensions('pdf')
		else:
			raise Exception('Unknown procedure')
		
		procedure.set_attribution(
			'Lee Yingtong Li (RunasSudo)',  # Author
			'Lee Yingtong Li (RunasSudo)',  # Copyright
			'Lee Yingtong Li (RunasSudo)'   # Year
		)
		return procedure

Gimp.main(FileJBIG2PDF.__gtype__, sys.argv)