#!/usr/bin/python
import calendar,math,md5,optparse,os,re,signal,struct,time,zlib,zipfile,shutil,sys
from romlib import *

#----------------------------------------------------------
# Rom sorting
#----------------------------------------------------------
class Dat:
	def __init__(self,name):
		self.name=name
		self.path=name.replace(' ','_')
		self.nbr=0
		self.found=0
		self.games=[]
class Game:
	def __init__(self,dat,name):
		self.dat=dat
		self.dat.games.append(self)
		self.name=name
		self.roms=[]
		self.found=0
		self.vdir=VDir(self.dat.path)
class Rom:
	def __init__(self,game,name,size,crc,md5):
		self.game=game
		self.game.roms.append(self)
		self.name=name
		self.size=size
		self.crc=crc
		self.md5=md5
		self.found=0
	def action_found(self,vfile):
		if self.found:
			print "  dupe"
		else:
			print "  added to",self.game.vdir.name
			self.game.vdir.add(vfile,self.name)
			self.found=1
			self.game.found+=1
			self.game.dat.found+=1
class RomDB:
	def __init__(self):
		self.dats=[]
		self.count=0
		self.fsize={}
		self.mode_delete=0
		self.printlast=0

	def printline(self,s):
		sys.stdout.write("\r"+" "*self.printlast+"\r"+s)
		sys.stdout.flush()
		self.printlast=len(s)

	def add_cmdat(self,vfile):
		allfile=vfile.read()
		if allfile.startswith("clrmamepro"):
			m=re.search('\tcategory "([^"]+)"',allfile)
			if m:
				dat=Dat(m.group(1))
			else:
				m=re.search(r'clrmamepro \(\s+name "([^"]+)"',allfile)
				if m:
					dat=Dat(m.group(1))
				else:
					raise "dat:Name not found"
			self.dats.append(dat)
			lm=re.findall(re.compile(r'game \(\s+name "([^"]+)"\s+description "([^"]+)"\s+(\trom.*?)^\)',re.S|re.M),allfile)
			del allfile
			dat.nbr=len(lm)
			romre=re.compile(r'\(\s+name "([^"]+)" size (\d+) crc ([A-Fa-z0-9]+) md5 ([A-Fa-z0-9]+)')
			for game in lm:
				g=Game(dat,game[0])
				l=game[2].split("\trom")
				for rom in l[1:]:
					m=re.search(romre,rom)
					if m:
						self.count+=1
						if self.count%512==1:
							self.printline("Loading: [%06d] %s"%(self.count,vfile.name))
						r=Rom(g,m.group(1),int(m.group(2)),m.group(3),m.group(4))
						if not self.fsize.has_key(r.size):
							self.fsize[r.size]={}
						self.fsize[r.size][r.md5]=r

	def add_cmdat_dir(self,d):
		for i in [f for f in VDir(d).listrec() if f.name.lower().endswith(".dat")]:
			self.add_cmdat(i)
		self.printline("Loaded: %d dats %d files"%(len(self.dats),self.count))
		print

	def set_mode_zip(self):
		for dat in self.dats:
			for game in dat.games:
				zipname=game.name+".zip"
				game.vdir=VZipDir(os.path.join(game.dat.path,zipname).replace(' ','_'))

	def set_mode_zip_merged(self):
		for dat in self.dats:
			for game in dat.games:
				n=game.name
#				Not for multifiles archives
				if n.find('(')!=-1 and len(game.roms)==1:
					n=n[:n.find('(')].strip()
				zipname=n+".zip"
				game.vdir=VZipDir(os.path.join(game.dat.path,zipname).replace(' ','_'))

	def set_mode_delete(self):
		self.mode_delete=1

	def scan(self,d):
		for dat in self.dats:
			vd=VDir(os.path.join(d,dat.path))
			found={}
			for vfile in vd.listrec():
				found[vfile.name]=1
			for game in dat.games:
				for rom in game.roms:
					if found.has_key(rom.name):
						rom.found=1
						rom.game.found+=1
						rom.game.dat.found+=1
			print "Scanned: %s (%d/%d)"%(dat.name,dat.found,dat.nbr)

	def rebuild(self,d):
		tmp=VDir(d).listrec()
		for vfile in tmp:
			dmd5=self.fsize.get(vfile.size,0)
			if dmd5!=0:
				fb=vfile.read()
				if len(fb):
					fmd5=''.join(['%02x'%ord(i) for i in md5.new(fb).digest()])
					r=dmd5.get(fmd5,0)
					if r!=0:
#						print "%s:\n  match %s/%s/%s"%(vfile.fullname,r.game.dat.name,r.game.name,r.name)
						print "%s:\n  match %s/%s"%(vfile.fullname,r.game.dat.name,r.name)
						r.action_found(vfile)
						# mode unlink
						if self.mode_delete:
							vfile.unlink()
		for dat in self.dats:
			print "Rebuilt: %s (%d/%d)"%(dat.name,dat.found,dat.nbr)

	def solidify_run(self,c):
		print c
		rc=os.spawnvp(os.P_WAIT, c[0], c)
		if rc:
			print "Err",c
			sys.exit()

	def solidify(self,d):
		for dat in self.dats:
			if os.path.isdir(dat.path):
				for i in [os.path.join(dat.path,i) for i in os.listdir(dat.path)]:
					if i.endswith(".zip") or os.path.isdir(i):
						base=i
						if i.endswith(".zip"):
							base=i[:-4]
						f7z=base+".7z"
						if os.path.isfile(f7z):
							self.solidify_run(['7za','x','-o'+base,f7z])
							os.unlink(f7z)
						fzip=base+".zip"
						if os.path.isfile(fzip):
							self.solidify_run(['7za','x','-o'+base,i])
							os.unlink(i)
						cwd=os.getcwd()
						os.chdir(base)
#						self.solidify_run(['7za','a','-mx=9'       ,os.path.join('..',os.path.basename(f7z)),'.'])
						self.solidify_run(['7za','a','-m0=LZMA:d23',os.path.join('..',os.path.basename(f7z)),'.'])
						os.chdir(cwd)
						shutil.rmtree(base)

def uncompress_all_walk(d,dirname,names):
	for fn in [os.path.join(dirname,i) for i in names]:
		if os.path.isfile(fn):
#			if fn.lower().endswith(".gz"):
#				c=['gzip','-d',fn]
#				print c
#				rc=os.spawnvp(os.P_WAIT, c[0], c)
			ext=['zip','7z','gz','bz2','tar','lzh','lha','arj','cab','rar']
			for i in ext:
				if fn.lower().endswith("."+i):
					dest=os.path.dirname(fn)
					c=['7z','x','-o'+dest,'-aou',fn]
					print c
					rc=os.spawnvp(os.P_WAIT, c[0], c)
					if rc==0:
						os.unlink(fn)
					else:
						if fn.lower().endswith(".rar"):
							print "Warning: p7zip-rar might be absent"
						if fn.lower().endswith(".lha") or fn.lower().endswith(".lzh"):
							c=['lha','xfw='+dest,fn]
							print c
							rc=os.spawnvp(os.P_WAIT, c[0], c)
							if rc==0:
								os.unlink(fn)
						else:
							print "Error while processing",fn

def uncompress_all(d):
	if os.path.isdir(d):
		os.path.walk(d,uncompress_all_walk,d)

#----------------------------------------------------------
# Automatic player
#----------------------------------------------------------

def main():
	parser = optparse.OptionParser()
	# modes
	parser.add_option("-r", "--rebuild", action="store_true", dest="rebuild", default=False, help="run in rebuild mode (default)")
	parser.add_option("-u", "--unzip", action="store_true", dest="unzip", default=False, help="pre-unzip all files mode")
	# options
	parser.add_option("-d", "--dats", dest="dats", default="dats", metavar="FILE", help="directory containings CM dats files (default: dats)")
	parser.add_option("-s", "--sort", dest="sort", default="sort", metavar="FILE", help="directory containings files to sort (default: sort)")
	parser.add_option("-o", "--output", dest="output", default=".", metavar="FILE", help="output directory (default: .)")
	# flags
	parser.add_option("-k", "--keep", dest="delete", action="store_false", default=True, help="keep matching files in sort (default: no)")
	parser.add_option("-1", "--nomerge", dest="merge", action="store_false", default=True, help="don't merge output zip (default: no)")
	parser.add_option("-7", "--solid", dest="solid", action="store_true", default=False, help="solidify to .7z archive (default: no)")

	(o, a) = parser.parse_args()
	if o.unzip:
		uncompress_all(o.sort)
	else:
#		signal.signal(signal.SIGCHLD, signal.SIG_IGN)
		rdb=RomDB()
		rdb.add_cmdat_dir(o.dats)
		if o.delete:
			rdb.set_mode_delete()
		if o.merge:
			rdb.set_mode_zip_merged()
		rdb.scan(o.output)
		rdb.rebuild(o.sort)
		if o.solid:
			rdb.solidify(o.sort)

main()