package main

import (
	"flag"
	"io"
	"os"

	"bitbucket.org/ausocean/av/stream/mts"
	"github.com/Comcast/gots/packet"
)

const (
	errBadInPath     = "No file path provided, or file does not exist"
	errCantCreateOut = "Can't create output file"
	errCantGetPid    = "Can't get pid from packet"
	errReadFail      = "Read failed"
	errWriteFail     = "Write to file failed"
	usage            = "The path to the file to be repaired"
)

var ccMap = map[int]byte{
	mts.PatPid:   1,
	mts.PmtPid:   1,
	mts.VideoPid: 0,
}

type Packet [mts.PacketSize]byte

func (p *Packet) setCC(cc byte) {
	(*p)[3] |= cc & 0xf
}

func main() {
	// Deal with input flags
	inPtr := flag.String("path", "", usage)
	outPtr := flag.String("out", "out.ts", usage)
	flag.Parse()

	// Try and open the given input file, otherwise panic - we can't do anything
	inFile, err := os.Open(*inPtr)
	defer inFile.Close()
	if err != nil {
		panic(errBadInPath)
	}

	// Try and create output file, otherwise panic - we can't do anything
	outFile, err := os.Create(*outPtr)
	defer outFile.Close()
	if err != nil {
		panic(errCantCreateOut)
	}

	// Read each packet from the input file reader
	var p Packet
	for {
		// If we get an end of file then return, otherwise we panic - can't do anything else
		if _, err := inFile.Read(p[:mts.PacketSize]); err == io.EOF {
			return
		} else if err != nil {
			panic(errReadFail + ": " + err.Error())
		}

		// Get the pid from the packet and set the cc based on this pid using our map
		pid, err := packet.Pid((*packet.Packet)(&p))
		if err != nil {
			panic(errCantGetPid)
		}
		p.setCC(ccFor(int(pid)))

		// Write this packet to the output file
		if _, err := outFile.Write(p[:]); err != nil {
			panic(errWriteFail + ": " + err.Error())
		}
	}
}

// ccFor gets the next cc for the given pid
func ccFor(pid int) byte {
	cc := ccMap[pid]
	ccMap[pid] = (cc + 1) & 0xf
	return cc
}