import sys, re, os
from position import *

def error(s, ec = 1):
	sys.stderr.write("ERROR: %s\n" % s)
	sys.exit(ec)

re_keyVal = re.compile(r'([^:]*\S)\s*:\s*(.*)\r?\n?')
re_idxEntry = re.compile(r'([IPB])\s+([0-9A-F]+)\s+([0-9A-F]+)(?:\s+([0-9a-f]+)\s+([0-9]+))?$')

def splitKeyVal(s):
	ma = re_keyVal.match(s)
	if ma:
		return ma.group(1), ma.group(2)
	else:
		return None

def isEDL(filename):
	return file(filename).read(8) == "ADMW0002"

class EDL:
	class Segment:
		def __init__(self, parent, start, size, ref):
			self._parent = parent
			self._start = start
			self._size = size
			self._ref = ref

		def startFrame(self):
			return self._start

		def endFrame(self):
			return self._start + self._size

		def frameCount(self):
			return self._size

		def length(self):
			return self._size / self.openIDX().fps()

		def videoIndex(self):
			return self._ref

		def videoFilename(self):
			return self._parent.videos[self._ref].filename

		def openIDX(self):
			return self._parent.openIDX(self._ref)

		def bounds(self):
			return self.startFrame(), self.endFrame()

		def _nextIntraFrame(self, idx, start, offset):
			if not (start >= 0 and start < len(idx)):
				raise IndexError("illegal frame %d" % start)

			if offset == None:
				nif = self._nextIntraFrame(idx, start, 1)
				pif = self._nextIntraFrame(idx, start, -1)
				if nif == None:
					return pif
				if pif == None:
					return nif
				if nif - start < start - pif:
					return nif
				else:
					return pif

			result = start
			while idx[result][0] != "I":
				result += offset
				if not (result >= 0 and result < len(idx)):
					return None

			return result

		def intraFrameBounds(self, mode = "open"):
			idx = self.openIDX()

			sOfs, eOfs = {
				"open" : (-1, 1),
				"close" : (1, -1),
				"nearest" : (None, None),
				}[mode]

			sf = self._nextIntraFrame(idx, self.startFrame(), sOfs)
			if sf == None:
				sf = self._nextIntraFrame(idx, self.startFrame(), -sOfs)

			ef = self._nextIntraFrame(idx, self.endFrame(), eOfs)
			if ef == None:
				ef = self._nextIntraFrame(idx, self.endFrame(), -eOfs)

			return sf, ef

	class Video:
		def __init__(self, filename):
			self.filename = filename
			self.idx = None

	def __init__(self, filename):
		edl = file(filename).readlines()
		if not edl[1].endswith(" videos\n"):
			error("can't interpret # of videos! (%s)" % edl[1].strip())
		videoCount = int(edl[1][:-8])
		self.videos = []
		for i in range(videoCount):
			if not edl[2+i].startswith("Name : "):
				error("could not find video name: %s" % edl[2+i].strip())
			self.videos.append(self.Video(edl[2+i][7:].strip()))
		if not edl[2+videoCount].endswith(" segments\n"):
			error("could not find number of segments: %s" % edl[2+videoCount].strip())
		segmentCount = int(edl[2+videoCount][:-10])
		self.segments = []
		for i in range(segmentCount):
			for j in range(3):
				k, v = splitKeyVal(edl[3+videoCount+3*i+j])
				if k == "Start":
					start = int(v)
				if k == "Size":
					size = int(v)
				if k == "Ref":
					ref = int(v)
			self.segments.append(self.Segment(self, start, size, ref))

	def __len__(self):
		return len(self.segments)

	def __getitem__(self, index):
		return self.segments[index]

	def totalFrameCount(self):
		return sum([s.frameCount() for s in self])

	def totalLength(self):
		return sum([s.length() for s in self])

	def openIDX(self, ref):
		video = self.videos[ref]
		if not video.idx:
			if not video.filename.endswith(".idx"):
				error("openIDX: expected MPEG .idx file (got %s)!"
					  % (video.filename, ))
			try:
				video.idx = IDX(video.filename)
			except IOError, e:
				if e.errno == 2 and os.path.isabs(video.filename):
					dir, base = os.path.split(video.filename)
					video.idx = IDX(base)
		return video.idx

	def dump(self, outFile,
			 verbose = False, blockSize = 128*1024):
		for s in self:
			idx = s.openIDX()
			# blind hack: dump valid PES stream...
			someOffset = idx[0][2]-idx[0][1]
			#someOffset = 0x3c-0xe

			sf, ef = s.intraFrameBounds()
# 			if s.endFrame() < ef and idx[s.endFrame()-1][0] == "I":
# 				sys.stderr.write(
# 					"end frame is 1 after I-frame, using anyways..\n")
# 				ef = s.endFrame()
			startPos = idx[sf][2] - someOffset
			endPos = idx[ef][2] - someOffset

			bytesToCopy = endPos - startPos
			if verbose:
				sys.stderr.write("copying %d - %d = %d bytes\n" % (
					startPos, endPos, bytesToCopy))
			inFile = file(idx.video)
			inFile.seek(startPos)
			while bytesToCopy > 0:
				buf = inFile.read(min(blockSize, bytesToCopy))
				outFile.write(buf)
				bytesToCopy -= len(buf)
				if verbose:
					sys.stderr.write("(%d to go) \r" % (bytesToCopy, ))
					sys.stderr.flush()

class IDX:
	def __init__(self, filename):
		self.idxFile = file(filename).readlines()
		self.video = self.idxFile[3].strip()
		self.info = [int(field) for field in
					 self.idxFile[4].strip().split(" ")]

	def __len__(self):
		return len(self.idxFile)-5

	def frameWidth(self):
		return self.info[0]

	def frameHeight(self):
		return self.info[1]

	def fps(self):
		return self.info[2] / 1000.0

	def __getitem__(self, index):
		rawItem = self.idxFile[5+index]
		ma = re_idxEntry.match(rawItem)
		if ma:
			frameType = ma.group(1)
			streamPos = int(ma.group(2), 16) # meaning unsure
			if frameType == "I":
				filePos = int(ma.group(3), 16)
				return (frameType, streamPos, filePos, ma.group(4), ma.group(5))
			else:
				size = int(ma.group(3))
				return (frameType, streamPos, size)
		else:
			error("'%s' did not match IDX regexp!" % rawItem.strip())

def convertEDL(edl, includeFrameCount = True, convert = None):
	result = []
	for s in edl:
		idx = s.openIDX()
		filename = idx.video
		if convert == "absolute":
			filename = os.path.abspath(filename)
		if includeFrameCount:
			result.append("-sb %d -frames %d '%s'" %
						  (idx[s.startFrame()][2], s.frameCount(), filename))
		else:
			result.append("-sb %d '%s'" %
						  (idx[s.startFrame()][2], filename))
	return result

def debugEDL(edl):
	print "%d frames total (%s seconds)" % (
		edl.totalFrameCount(), edl.totalLength())

	position = Position()
	for i, s in enumerate(edl):
		position += s.length()
		print "segment %d ends at %s" % (i+1, position)

	for s in edl:
		sf, ef = s.intraFrameBounds()
		idx = s.openIDX()

		print "start frame %d is" % s.startFrame(),
		if s.startFrame() > sf:
			print "%d frames after I-frame" % (s.startFrame() - sf)
			print "  ", idx[s.startFrame()]
		else:
			print "an I-frame"

		print "end frame %d is" % s.endFrame(),
		if s.endFrame() < ef:
			print "%d frames before I-frame" % (ef - s.endFrame())
			print "  ", idx[s.endFrame()]
		else:
			print "an I-frame"

def checkPattern(index, pattern = "IBBPBBPBBPBB"):
	frameCount = len(index)
	frameIndex = 0
	patternIndex = 0
	while True:
		while frameIndex < frameCount and index[frameIndex][0] == pattern[patternIndex]:
			frameIndex += 1
			patternIndex = (patternIndex + 1) % len(pattern)
		if frameIndex >= frameCount:
			break
		print "expected frame #%d to be an %s-frame (%s[%s]%s), but found %s-frame" % (frameIndex + 1, pattern[patternIndex], pattern[:patternIndex], pattern[patternIndex], pattern[patternIndex+1:], index[frameIndex][0])
		skipping = 0
		while frameIndex < frameCount and index[frameIndex][0] != pattern[0]:
			frameIndex += 1
			skipping += 1
		print "skipped %d frames, re-beginning with %s-frame at frame #%d" % (skipping, pattern[0], frameIndex + 1)
		patternIndex = 0

if __name__ == "__main__":
	edl = EDL(sys.argv[1])
	if len(sys.argv) > 2:
		if sys.argv[2] == "-":
			edl.dump(sys.stdout)
		else:
			debugEDL(edl)
			edl.dump(file(sys.argv[2], "w"), verbose = True)
	else:
		debugEDL(edl)
		
		print "\n".join(convertEDL(edl))
		print "(give additional filename argument to dump segment contents)"