#!/usr/bin/python

import sys, os
import psycopg2

class Migrator:
	def __init__(self):
		try:
			self.mOldConn = psycopg2.connect("dbname=shemov user=shemov password=shemov")
			self.mNewConn = psycopg2.connect("dbname=shemov2 user=shemov password=shemov")
		except psycopg2.OperationalError as ex:
			print ex,
			sys.exit(1)
		self.mGenreMap = {}
		self.mNewGenreMap = {}
		self.mActorMap = {}
		self.mNewActorMap = {}
		self.mSeriesMap = {}
		self.mSeriesPartMap = {}
		self.reset()
		self.migrate()

	def reset(self):
		cur = self.mNewConn.cursor()
		cur.execute("DELETE FROM series")
		cur.execute("SELECT setval('series_iseries_id__seq', 1)")
		cur.execute("SELECT setval('genres_igenres_id__seq', 1)")
		cur.execute("SELECT setval('seriesparts_seriesparts_id__seq', 1)")
		cur.execute("SELECT setval('actors_iactors_id__seq', 1)")
		cur.execute("DELETE FROM genres")
		cur.execute("DELETE FROM actors")

	def migrate(self):
		movieQuery = "SELECT imovid, ttitle, bisize, igenreid, iquality, idvd, iseriesno, ipartno, cmd5sum, tfilename FROM movies"
		movieCursor = self.mOldConn.cursor()
		movieCursor.execute(movieQuery)
		for rec in movieCursor:
			(movid, title, size, genre, quality, dvd, seriesno, partno, md5, filename) = rec
			print "Migrating %s %s (part %s)" % (title, seriesno, partno)
			genreName = self.genreName(genre)
			actorIds = self.actorIds(movid)
			actorNames = self.actorNames(actorIds)
			self.newEntry(title, seriesno, filename, md5, size, dvd, 1, partno, genreName, actorNames, quality)
		self.covers()

	def covers(self):
		q1 = "SELECT covers.tfilename, covers.imovid, tcovertype, covers.cmd5sum, ttitle, iseriesno FROM covers, movies WHERE covers.imovid = movies.imovid"
		q2 = "SELECT iseriesparts_id from seriesparts, series WHERE series.tseries_name = %s AND seriesparts.iseriespart = %s AND series.iseries_id = seriesparts.iseries_id"
		q3 = "INSERT INTO files(iseriespart_id, tfilename, cmd5sum, bisize, sifiletype) VALUES(%s, %s, %s, %s, %s)"
		oldCur = self.mOldConn.cursor()
		oldCur.execute(q1)
		newCur = self.mNewConn.cursor()
		newCur2 = self.mNewConn.cursor()
		coverTypesMap = {
			'front'		: 2,
			'back'		: 3,
			'general'	: 4
		}
		for cover in oldCur:
			(filename, movid, tcovertype, cmd5sum, title, series) = cover
			print "Fixing cover %s" % (filename,)
			series = series == -1 and 1 or series
			coverType = coverTypesMap[tcovertype]
			newCur.execute(q2, (title, series))
			seriespartids = newCur.fetchone()
			seriespartid = -1
			if len(seriespartids) != 0:
				seriespartid = seriespartids[0]
			path = "/home/am/movs/vids/archive/%s/%s/%s" % (cmd5sum[0], cmd5sum[1], filename)
			size = -1
			try:
				size = os.stat(path).st_size
			except OSError:
				print "Could not stat %s" % path
				continue
			newCur2.execute(q3, (seriespartid, filename, cmd5sum, size, coverType))
			self.mNewConn.commit()

	def newEntry(self, title, partno, filename, md5, size, dvd, filetype, fileno, genreName, actorNames, quality):
		seriesid = self.newSeries(title)
		seriespartid = self.newSeriesPart(seriesid, partno) 
		self.newFile(seriespartid, filename, md5, size, dvd, 1, fileno, quality)
		genreId = self.genreId(genreName)
		newCur = self.mNewConn.cursor()
		if genreId != -1:
			q1 = "INSERT INTO seriesparts_genremap(iseriesparts_id, igenres_id) VALUES(%s, %s)"
			newCur.execute(q1, (seriespartid, genreId))
			self.mNewConn.commit()
		actorIds = self.newActorIds(actorNames)
		if len(actorIds):
			q2 = "INSERT INTO seriesparts_actormap(iseriesparts_id, iactors_id) VALUES(%s, %s)"
			for aId in actorIds:
				try:
					newCur.execute(q2, (seriespartid, aId))
				except psycopg2.Error:
					pass
			self.mNewConn.commit()

	def newActorIds(self, actorNames):
		retval = []
		cur = self.mNewConn.cursor()
		for actor in actorNames:
			if actor in self.mNewActorMap.keys():
				retval.append(self.mNewActorMap[actor])
			else:
				q1 = "INSERT INTO actors(tactorname) VALUES(%s)"
				cur.execute(q1, (actor,))
				cur.execute("SELECT currval('actors_iactors_id__seq')")
				newId = cur.fetchone()[0]
				self.mNewConn.commit()
				self.mNewActorMap[actor] = newId
				retval.append(newId)
		return retval

	def genreId(self, genreName):
		if genreName in self.mNewGenreMap.keys():
			return self.mNewGenreMap[genreName]
		q1 = "INSERT into genres(tgenrename) VALUES(%s)"
		cur = self.mNewConn.cursor()
		try:
			cur.execute(q1, (genreName,))
		except psycopg2.Error as ex:
			print "Error %s" % (ex.pgerror,)
			self.mNewConn.rollback()
			return -1
		cur.execute("SELECT currval('genres_igenres_id__seq')")
		retval = cur.fetchone()
		self.mNewConn.commit()
		self.mNewGenreMap[genreName] = retval
		return retval

	def newFile(self, seriespartid, filename, md5, size, dvd, filetype, fileno, quality):
		q1 = "INSERT INTO files(iseriespart_id, tfilename, cmd5sum, bisize, idvd, sifiletype, sifileno, siquality) VALUES(%s, %s, %s, %s, %s, %s, %s, %s)"
		cur = self.mNewConn.cursor()
		try:
			cur.execute(q1, (seriespartid, filename, md5, size, dvd, filetype, fileno, quality))
		except psycopg2.Error as ex:
			print "Error: %s" %(ex.pgerror, )
			self.mNewConn.rollback()
		self.mNewConn.commit()


	def newSeriesPart(self, seriesid, part):
		part = part == -1 and 1 or part
		q1 = "SELECT iseriesparts_id FROM seriesparts WHERE iseries_id = %s AND iseriespart = %s"
		cur = self.mNewConn.cursor()
		cur.execute(q1, (seriesid, part))
		spartids = cur.fetchall()
		retval = -1
		if len(spartids) == 0:
			q2 = "INSERT INTO seriesparts(iseriespart, iseries_id) VALUES(%s, %s)"
			cur.execute(q2, (part, seriesid))
			cur.execute("SELECT currval('seriesparts_seriesparts_id__seq')")
			retval = cur.fetchall()[0][0]
			self.mNewConn.commit()
		else:
			retval = spartids[0][0]
		return retval
		
	def newSeries(self, seriesName):
		if seriesName in self.mSeriesMap.keys():
			return self.mSeriesMap[seriesName]
		q1 = "INSERT INTO series(tseries_name) VALUES(%s)"
		cur = self.mNewConn.cursor()
		try:
			cur.execute(q1, (seriesName, ))
		except psycopg2.Error as ex:
			print "Error: %s" % (ex.pgerror, )
			self.mNewConn.rollback()
			return -1

		q2 = "SELECT currval('series_iseries_id__seq')"
		cur.execute(q2)
		retval = cur.fetchone()[0]
		self.mNewConn.commit()
		self.mSeriesMap[seriesName] = retval
		return retval
	
	def actorIds(self, movId):
		query = "SELECT iactorid from movieactormap WHERE imovid = %s"
		cur = self.mOldConn.cursor()
		cur.execute(query, (movId,))
		return cur.fetchall()

	def genreName(self, genreId):
		if genreId in self.mGenreMap.keys():
			return self.mGenreMap[genreId]
		query = "SELECT tgenrename FROM genre WHERE igenreid = %s"
		cur = self.mOldConn.cursor()
		cur.execute(query, (genreId, ))
		name = cur.fetchone()
		if name:
			self.mGenreMap[genreId] = name[0]
		return self.mGenreMap[genreId]

	def actorNames(self, actorIds):
		retval = []
		query = "SELECT tactorname FROM actor where iactorid = %s"
		cur = self.mOldConn.cursor()
		for aId in actorIds:
			if aId in self.mActorMap.keys():
				retval.append(self.mActorMap[aId])
			else:
				cur.execute(query, (aId, ))
				name = cur.fetchone()
				if name:
					self.mActorMap[aId] = name[0]
					retval.append(name[0])
		return retval

if __name__ == '__main__':
	t = Migrator()