#!/usr/bin/env python
# -*- coding: utf8 -*-

from __future__ import with_statement
from glob import glob
from PIL import Image, ImageColor
import string
import os

def i16(c):
	'''read big endian short'''
	return ord(c[1]) + (ord(c[0])<<8)

def i32(c):
	'''read big endian long'''
	return ord(c[3]) + (ord(c[2])<<8) + (ord(c[1])<<16) + (ord(c[0])<<24)

def hc2tc_int(w):
	'''convert High Color to True Color'''
	r = (w & 0xf800) << 8
	g = (w & 0x7e0) << 5
	b = (w & 0x1f) << 3
	return r|g|b

def hc2tc_str(w):
	'''convert High Color to True Color'''
	r = (w & 0xf800) >> 8
	g = (w & 0x7e0) >> 3
	b = (w & 0x1f) << 3
	return chr(r)+chr(g)+chr(b)

def basename(s):
	return os.path.basename(os.path.splitext(s)[0])

#if __name__=="__main__":
	#format = 'TRS'
	#format_description = 'True Color Sprites trs Format'

base_path = 'E:/Steam/SteamApps/common/Dominions4/data'
debug = 1
offsets = {
        base_path + '/monster.trs': 0,
		base_path + '/flag.trs': 0
}

print offsets.keys()



for trs in offsets.keys():
	if os.path.isdir(basename(trs)) == False:
		os.mkdir(basename(trs), 0755)
	with open(trs, "rb") as fp:
		# HEAD

		s = fp.read(12)
		if s[:4] != 'TCSF':
			raise SyntaxError, 'not a Truecolor sprites trs'


		numsprites	= i16(s[4:])
		version		= i16(s[6:])
		xres		= i16(s[8:])
		filler		= i16(s[10:])
		
		
		if debug:
			print "Reading " + trs
			print "Version %d, %d sprites at width %d" % (version, numsprites, xres)
			print '=' * len(trs)

		mode = 'RGB'

		# Read TRS index
		sprites = []
		unpacked = []
		for spr in xrange(numsprites):
			s = fp.read(12)
			xsize		= ord(s[0])
			ysize		= ord(s[1])
			unknown		= i16(s[2:])
			offset_unpacked	= i32(s[4:])
			offset_packed	= i32(s[8:])
			
			if offset_packed and offset_unpacked:
				raise SyntaxError, 'image data both raw and compressed'
			
			if xsize == 0:
				xsize = 256
			if ysize == 0:
				ysize = 256
			
			if debug:
				print "Sprite %d: %dx%d, unknown: %d, unpacked: %d, packed: %d" % (spr, xsize, ysize, unknown, offset_unpacked, offset_packed)

			if offset_unpacked != 0:
				unpacked.append((spr, xsize, ysize, offset_unpacked))

			if offset_packed != 0:
				sprites.append((spr, xsize, ysize, offset_packed))

		# Extract raw sprite data
		for sprnbr, xsize, ysize, offset in unpacked:
			#file = basename(trs)+"/"+string.zfill(offset, len(str(os.path.getsize(trs))))+".png"
			#file = basename(trs)+"/"+string.zfill(sprnbr, 3)+".png"
			# army stuff
			print "Saving", file
			file = basename(trs)+"/"+string.zfill(offsets[trs]+sprnbr, 4)+".png"
			open(file, 'a').close()					# touch it
			try:
				fp.seek(offset)
				data = ''
				for i in xrange(xsize*ysize):
					w = i16(fp.read(2))
					data += hc2tc_str(w)
				im = Image.frombuffer(mode, (xsize, ysize), data, 'raw', mode, 0, 1)
				im.save(file, "PNG")
			except:
				print "error at", trs, "sprite", sprnbr, "offset", offset

		# Extract packed sprite data
		for sprnbr, xsize, ysize, offset in sprites:
			#file = basename(trs)+"/"+string.zfill(offset, len(str(os.path.getsize(trs))))+".png"
			#file = basename(trs)+"/"+string.zfill(sprnbr, 3)+".png"
			# army stuff
			file = basename(trs)+"/"+string.zfill(offsets[trs]+sprnbr, 4)+".png"
			open(file, 'a').close()					# touch it
			if 1:
				fp.seek(offset)
				repeat = 0
				line_wrote = 0
				data = ''
				size = (xres, ysize)

				chunks = i16(fp.read(2))+1
				for i in xrange(chunks):
					scrptr_offs  = i16(fp.read(2))
					scrptr_offs /= 2			# we're still in High-Color mode ;)

					# screen offset
					if scrptr_offs > xres:
						scrptr_offs += line_wrote
						line_wrote = 0
					else:
						line_wrote += scrptr_offs
	
					data += '\0' * 3*scrptr_offs

					repeat = i16(fp.read(2))
					# check for negative sign
					if repeat & 0x8000:
						continue
					repeat += 1
					
					# pixel data
					for j in xrange(repeat):
						w = i16(fp.read(2))
						data += hc2tc_str(w)
					line_wrote += repeat

				# pad data
				if len(data) > 3*size[0]*size[1]:
					data = data[:3*size[0]*size[1]]
				else:
					data += '\0' * (3*size[0]*size[1]-len(data))
				
			try:
				im = Image.frombuffer(mode, size, data, 'raw', mode, 0, 1)
				
				# don't swap alpha channels, TO ACTIVATE REMOVE UNTIL "CONTINUE" STATEMENT
				#region = im.crop((0,0,xsize,ysize))
				#region.save(file, "PNG")
				#continue
				
				# add alpha channel
				black = (0, 0, 0)
				cyan = (248, 0, 248)
				white = (255, 255, 255)
				
				image = im.load()
				mask = Image.new("1", (xsize,ysize))
				maskD = mask.load()
				for y in range(ysize):
					for x in range(xsize):
						r,g,b = image[x,y]
						if (g == 0 and r > 245 and b > 245):
							# cyan, shadow it
							image[x,y] = black
							maskD[x,y] = 1
						elif image[x,y] != black:
							maskD[x,y] = 1
						else:
							# black, make it transparent
							image[x,y] = white
				
				region = im.crop((0,0,xsize,ysize))
				region.putalpha(mask)
				
				print "Saving", file
				
				region.save(file, "PNG")
			except:
				print "error at", trs, "sprite", sprnbr, "offset", offset
				
			print "Done!";
			print "";