#!/usr/bin/env python

import sys
sys.path.append('/Users/jgurney/p4/bktrau/info')

import itertools
import mpegts
import sys
import sets

def usage():
	print >>sys.stderr, 'Usage: %s <file> <pmtpid> <pid> ...' % sys.argv[0]
	sys.exit(1)

def genpats(pmt):
	BASEPAT = map(None, "G@\x00\x10\x00\x00\xb0\r\x00\x00\xc1\x00\x00\x00\x01\xe0\x00\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff")

	patidx = 4 + 1	# TS header + pointer table
	BASEPAT[patidx + 10] = chr(0xe0 | ((pmt >> 8) & 0x1f))
	BASEPAT[patidx + 11] = chr(pmt & 0xff)
	newcrc = mpegts.psip_calc_crc32(''.join(BASEPAT[patidx:patidx + 12]))
	newcrc = map(lambda x, crc = newcrc: chr((crc >> (8 * (3 - x))) & 0xff), range(4))
	BASEPAT[patidx + 12:patidx + 16] = newcrc

	assert len(BASEPAT) == mpegts.TSPKTLEN

	ret = []
	old = ord(BASEPAT[3]) & 0xf0
	for i in range(16):	# continuity
		BASEPAT[3] = chr(old | i)
		ret.append(''.join(BASEPAT))

	return ret

def producets(inp, pmtpid, *pids):
	print `inp`, `pmtpid`, `pids`
	pats = itertools.cycle(genpats(pmtpid))
	# XXX - check if all pids are ints? in range?
	pids = sets.Set(pids)

	stream = mpegts.TSPStream(inp)
	for i in stream:
		frst = ord(i[1])
		# Get first and error bits for testing.
		pid = (frst & 0x1f) << 8 | ord(i[2])
		if pid == pmtpid and (frst & 0xc0) == 0x40:
			yield pats.next()
			# XXX - we probably want to rewrite the PMT to only
			# include the pids we are sending.
			yield i
		elif pid in pids:
			yield i

def main():
	if len(sys.argv) < 3:
		usage()

	pmtpid = int(sys.argv[2])
	pids = map(int, sys.argv[3:])

	inp = open(sys.argv[1])
	out = sys.stdout

	producer = producets(inp, pmtpid, *pids)
	filter(lambda x: out.write(x), producer)

if __name__ == '__main__':
	main()