import bpy
import math
import struct
import os
import time
import tempfile

bl_info = {
	'name': 'GRP Tools',
	'description': 'Output StarCraft GRP files',
	'author': 'John Wharton',
	'version': (1, 0, 0),
	'blender': (2, 80, 0),
	'location': 'Output > GRP Tools',
	'warning': 'Requires StarCraft Python Modding Suite and Pillow',
	'doc_url': 'https://gitlab.com/Solstice245/grp-tools',
	'tracker_url': 'https://gitlab.com/Solstice245/grp-tools/-/issues',
	'category': 'Scene',
}


DESC_GRP_PATH = '''
	Filepath of the output GRP file (relative to MPQ archive directory)
'''
DESC_GRP_PAL = '''
	Optionally specifies which palette to quantize the render output with.\n\n
	Value should match the name of an .act file in _PMS\\Palettes
'''
DESC_GRP_OUTBOOL = '''
	Determines whether this GRP will be rendered and output when the "Output GRPs" operation is executed
'''
DESC_GRP_SIZE = '''
	If set to a non-zero value, this value will be used to determine the resolution of the GRP instead of the scene values
'''
DESC_GRP_FRAMESTART = '''The frame which the renderer will start at for this GRP'''
DESC_GRP_FRAMENED = '''The frame which the renderer will end at for this GRP'''

ENUM_CAM_ANGLES = (
	('STATIC', 'Static', 'Renders 1 angle per frame (generally only used by structures)'),
	('RIGHT', 'Right', 'Renders 17 angles per frame, which the engine uses to simulate 32 directional movement'),
	('OMNI', 'Omni', 'Renders 32 angles per frame'),
)


class GRP_CONFIG(bpy.types.PropertyGroup):
	output_path: bpy.props.StringProperty(options=set(), default='unit\\spritename', description=DESC_GRP_PATH)
	output_bool: bpy.props.BoolProperty(options=set(), default=True, description=DESC_GRP_OUTBOOL)
	color_palette: bpy.props.StringProperty(options=set(), default='!gamebasic', description=DESC_GRP_PAL)
	render_size: bpy.props.IntProperty(options=set(), min=0, max=256, default=0, description=DESC_GRP_SIZE)
	camera_angles: bpy.props.EnumProperty(options=set(), items=ENUM_CAM_ANGLES, default='OMNI')
	frame_start: bpy.props.IntProperty(options=set(), default=0, description=DESC_GRP_FRAMESTART)
	frame_end: bpy.props.IntProperty(options=set(), default=0, description=DESC_GRP_FRAMENED)


class UI_UL_grp_config(bpy.types.UIList):
	bl_idname = 'UI_UL_grpt_grp_config'

	def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index, flt_flag):
		if self.layout_type in {'DEFAULT', 'COMPACT'}:
			row = layout.row()
			row.prop(item, 'output_path', text='', emboss=False, icon_value=icon)
			row.prop(item, 'output_bool', text='', emboss=False, icon='CHECKBOX_HLT' if item.output_bool else 'CHECKBOX_DEHLT')


class SCENE_PT_grpt(bpy.types.Panel):
	bl_idname = 'SCENE_PT_grpt'
	bl_label = 'GRP Tools'
	bl_space_type = 'PROPERTIES'
	bl_region_type = 'WINDOW'
	bl_context = 'output'
	bl_options = set()

	@classmethod
	def poll(cls, context):
		return context.scene

	def draw(self, context):
		layout = self.layout
		layout.use_property_split = True
		scene = context.scene

		layout.operator('grpt.grp_output', icon='OUTPUT')
		row = layout.row()
		row.use_property_split = False
		row.prop(scene, 'grpt_output_preview', text='Preview Output')
		sub = row.column()
		sub.active = scene.grpt_output_preview
		sub.prop(scene, 'grpt_preview_only', text='Preview Only')
		layout.separator()
		uilist_rows = 5 if len(scene.grpt_grp_config) else 3
		row = layout.row()

		row.template_list(
			'UI_UL_grpt_grp_config', 'grpt_grp_config', scene,
			'grpt_grp_config', scene, 'grpt_grp_output_index',
			rows=uilist_rows,
		)

		col = row.column(align=True)
		col.operator('grpt.grp_config_add', text='', icon='ADD')
		col.operator('grpt.grp_config_del', text='', icon='REMOVE')
		col.separator()
		col.operator('grpt.grp_config_mov', text='', icon='TRIA_UP').shift = -1
		col.operator('grpt.grp_config_mov', text='', icon='TRIA_DOWN').shift = 1

		if not len(scene.grpt_grp_config) or scene.grpt_grp_output_index < 0:
			return

		grp_config = scene.grpt_grp_config[scene.grpt_grp_output_index]
		row = layout.row(align=True)
		row.prop(grp_config, 'frame_start', text='Frame Range')
		row.prop(grp_config, 'frame_end', text='')
		layout.prop(grp_config, 'render_size', text='Resolution')
		layout.prop(grp_config, 'color_palette', text='Color Palette')
		layout.prop(grp_config, 'camera_angles', text='Camera Angles')
		layout.separator()
		calc_frames = grp_config.frame_end - grp_config.frame_start + 1
		calc_angles = 1 if grp_config.camera_angles == 'STATIC' else 17 if grp_config.camera_angles == 'RIGHT' else 32
		layout.label(text=f'GRP Output Image Count: {calc_frames * calc_angles}')


class SCENE_OT_grp_config_add(bpy.types.Operator):
	bl_idname = 'grpt.grp_config_add'
	bl_label = 'Add GRP Item'
	bl_description = 'Adds a new item to the collection'

	def invoke(self, context, event):
		scene = context.scene
		scene.grpt_grp_config.add()
		scene.grpt_grp_output_index = len(scene.grpt_grp_config) - 1
		return {'FINISHED'}


class SCENE_OT_grp_config_del(bpy.types.Operator):
	bl_idname = 'grpt.grp_config_del'
	bl_label = 'Remove GRP Item'
	bl_description = 'Removes the active item from the collection'

	def invoke(self, context, event):
		scene = context.scene

		if scene.grpt_grp_output_index not in range(len(scene.grpt_grp_config)):
			return {'FINISHED'}

		scene.grpt_grp_config.remove(scene.grpt_grp_output_index)
		shift = 1 if scene.grpt_grp_output_index == len(scene.grpt_grp_config) else 0
		scene.grpt_grp_output_index = scene.grpt_grp_output_index - shift

		return {'FINISHED'}


class SCENE_OT_grp_config_mov(bpy.types.Operator):
	bl_idname = 'grpt.grp_config_mov'
	bl_label = 'Move GRP Item'
	bl_description = 'Moves the active item up/down in the collection'

	shift: bpy.props.IntProperty()

	def invoke(self, context, event):
		scene = context.scene

		if scene.grpt_grp_output_index not in range(len(scene.grpt_grp_config)):
			return {'FINISHED'}

		collection.move(scene.grpt_grp_output_index, scene.grpt_grp_output_index + self.shift)
		scene.grpt_grp_output_index += self.shift

		return {'FINISHED'}


class SCENE_OT_grp_output(bpy.types.Operator):
	bl_idname = 'grpt.grp_output'
	bl_label = 'Output GRPs'
	bl_description = 'Renders the scene for all GRP entries marked for output within their respective frame ranges'
	bl_options = {'REGISTER'}

	@classmethod
	def poll(cls, context):
		return context.scene

	def execute(self, context):
		scene = context.scene

		if not len(scene.grpt_grp_config):
			return {'FINISHED'}

		render_dir = tempfile.gettempdir()

		t_exec = time.time()

		# store/overwrite user defined render settings
		user_frame = scene.frame_current
		user_file_path = scene.render.filepath
		user_file_format = scene.render.image_settings.file_format
		user_color_depth = scene.render.image_settings.color_depth
		user_compression = scene.render.image_settings.compression

		scene.render.image_settings.file_format = 'PNG'
		scene.render.image_settings.color_depth = '8'  # higher depths not needed
		scene.render.image_settings.compression = 0  # maximize write speed of PNGs

		for grp_config in scene.grpt_grp_config:
			if not grp_config.output_bool:
				continue

			t_grp = time.time()

			sprite_size = grp_config.render_size if grp_config.render_size else scene.render.resolution_x
			assert sprite_size <= 256, f'GRP {grp_config.output_path} sprite size is out of range(0, 256)'

			if grp_config.camera_angles == 'OMNI':
				num_angles = 32
				rot_angle = 360 / 32
			elif grp_config.camera_angles == 'RIGHT':
				num_angles = 17
				rot_angle = 180 / 16
			else:  # static
				num_angles = 1
				rot_angle = 0

			frames = list(range(grp_config.frame_start, grp_config.frame_end + 1))
			frame_img_names = []

			if not len(frames):
				continue

			print('Output GRP:', grp_config.output_path)

			if scene.camera.parent:
				camera_rotator = scene.camera.parent
			else:
				camera_rotator = bpy.data.objects.new('Rotator', None)
				bpy.context.scene.collection.objects.link(camera_rotator)
				scene.camera.parent = camera_rotator

			user_rotator_mode = camera_rotator.rotation_mode

			if camera_rotator.rotation_mode in ('QUATERNION', 'AXIS_ANGLE'):
				camera_rotator.rotation_mode = 'XYZ'

			for f in frames:
				scene.frame_set(f)
				print('Output frame:', scene.frame_current)

				frame_number = str(f - grp_config.frame_start).zfill(3)  # adds leading zeroes up to length 3

				# store user defined camera rotation angle
				user_camera_ob_angle = camera_rotator.rotation_euler.z

				if num_angles != 1:
					frame_time = time.time()

				for ii in range(num_angles):

					if num_angles != 1:
						angle = ii * rot_angle + 180
						camera_rotator.rotation_euler.z = math.radians(angle)

					frame_filename = f'{frame_number}.{str(ii).zfill(2)}{os.extsep}png'
					scene.render.filepath = os.path.join(render_dir, frame_filename)
					frame_img_names.append(scene.render.filepath)
					bpy.ops.render.render(animation=False, use_viewport=False, write_still=True)

				if num_angles != 1:
					print('Output frame', scene.frame_current, 'time:', time.time() - frame_time)

				# restore user defined camera rotation angle
				camera_rotator.rotation_euler.z = user_camera_ob_angle

			camera_rotator.rotation_mode = user_rotator_mode

		# restore user defined render settings
		scene.frame_set(user_frame)
		scene.render.filepath = user_file_path
		scene.render.image_settings.file_format = user_file_format
		scene.render.image_settings.color_depth = user_color_depth
		scene.render.image_settings.compression = user_compression

		print('Total output time:', time.time() - t_exec)

		return {'FINISHED'}


classes = (
	GRP_CONFIG,
	UI_UL_grp_config,
	SCENE_PT_grpt,
	SCENE_OT_grp_config_add,
	SCENE_OT_grp_config_del,
	SCENE_OT_grp_config_mov,
	SCENE_OT_grp_output,
)


DESC_GRPT_PREVIEW = '''Open output images with the system default image previewing program as they are completed'''
DESC_GRPT_PREVIEWONLY = '''While "Preview Output" is on, this option disables the actual GRP file output'''


def register():
	for cls in classes:
		bpy.utils.register_class(cls)

	bpy.types.Scene.grpt_output_preview = bpy.props.BoolProperty(options=set(), description=DESC_GRPT_PREVIEW)
	bpy.types.Scene.grpt_preview_only = bpy.props.BoolProperty(options=set(), description=DESC_GRPT_PREVIEWONLY)
	bpy.types.Scene.grpt_grp_config = bpy.props.CollectionProperty(type=GRP_CONFIG)
	bpy.types.Scene.grpt_grp_output_index = bpy.props.IntProperty(options=set(), default=-1)


def unregister():
	for cls in reversed(classes):
		bpy.utils.unregister_class(cls)


if __name__ == '__main__':
	register()