import os
import zipfile
import time

def is_ztd(str):
	if any([x in str for x in zt_excludes]):
		return
	return str.lower().endswith(".ztd")

# the ZT root folder
zt_dir = "C:/Program Files (x86)/Microsoft Games/Zoo Tycoon"
# also try these subfolders
zt_subfolders = ("Updates", "XPACK1", "XPACK2", "ZUPDATE", "ZUPDATE1", "dlupdate", "dupdate")
# skip ZTD files whose names include any of the following strings
zt_excludes = ("ztatb","zts")

def query_entity(name, file_to_ztd, lowercase_to_original_name):
	print("Getting all files for "+name)
	lower_name = name.lower()
	# get all content files in their original name formatting
	content_files = [lowercase_to_original_name[n] for n in lowercase_to_original_name if lower_name in n]
	# for each of these, get the ZTD it should be extracted from
	ztd_paths = [file_to_ztd[n] for n in content_files]

	# now open each of the required ZTD files once
	unique_ztd_paths = set(ztd_paths)
	ztd_path_to_ztd_file = {}
	for ztd_path in unique_ztd_paths:
		ztd_path_to_ztd_file[ztd_path] = zipfile.ZipFile(ztd_path)

	# just unzip each content file from its ZTD into the CWD
	# could be made much faster by unzipping all required files per ZTD in bulk
	for content, ztd in zip(content_files, ztd_paths):
		print("extract "+content+" from "+ztd)
		try:
			ztd_path_to_ztd_file[ztd].extract(content)
		except:
			print("ERROR: Could not unzip "+file)

def read_ztd_lib():
	# take the start time
	start = time.clock()

	#this dict remembers all files and where an updated version was found so it can be fetched to construct the animal
	file_to_ztd = {}
	#lowercase for lookup
	lowercase_to_original_name = {}

	# create a list of all folders that might contain ZTD files
	folders = [zt_dir,] + [os.path.join(zt_dir, sub) for sub in zt_subfolders]
	folders = [f for f in folders if os.path.isdir(f)]

	# get all ZTD files from those folders and store their absolute file paths in a list
	# note: should make sure the folders are sorted in the order they are read by the game, so overwriting works as intended
	ztd_files = [os.path.join(zt_dir, ztd_name) for ztd_name in os.listdir(zt_dir) if is_ztd(ztd_name)]

	# go over all ZTD files
	for ztd_path in ztd_files:
		# see if it is a ZTD with ZIP compression
		try:
			ztd_file = zipfile.ZipFile(ztd_path)
		except zipfile.BadZipFile:
			print(ztd_path+' is a not a ZIP file! Maybe a RAR file?')
			continue
		# just query the ZTD's entries
		for content_raw in ztd_file.namelist():
			# just read original ai files for now
			# if content_raw.endswith(".ai"):
				# normalize the path to make sure new versions overwrite previous entries
				content_lower = os.path.normpath(content_raw.lower())
				# link to where this file was found; overwrites older entry so we get the latest version
				file_to_ztd[content_raw] = ztd_path
				# store the link from this idealized file name back to the actual one
				lowercase_to_original_name[content_lower] = content_raw
	duration = time.clock()-start
	# print all files that
	print(file_to_ztd.keys())
	print("Done in {0:.2f} seconds".format(duration))
	return file_to_ztd, lowercase_to_original_name

file_to_ztd, lowercase_to_original_name = read_ztd_lib()
query_entity("lodgpine", file_to_ztd, lowercase_to_original_name)
query_entity("graywolf", file_to_ztd, lowercase_to_original_name)