av/cmd/revid-cli/main.go

/*
NAME
  revid-cli - command line interface for Revid.

DESCRIPTION
  See Readme.md

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>
  Jack Richardson <jack@ausocean.org>

LICENSE
  revid-cli is Copyright (C) 2017-2018 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
*/

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime/pprof"
	"strconv"
	"time"
	"unicode"

	"bitbucket.org/ausocean/av/revid"
	"bitbucket.org/ausocean/iot/pi/netsender"
	"bitbucket.org/ausocean/utils/smartlogger"
)

const (
	// progName is the program name for logging purposes.
	progName = "revid-cli"

	// Logging is set to INFO level.
	defaultLogVerbosity = smartlogger.Info
)

// Indexes for configFlags
const (
	inputPtr = iota
	inputCodecPtr
	outputPtr
	rtmpMethodPtr
	packetizationPtr
	quantizationModePtr
	verbosityPtr
	framesPerClipPtr
	rtmpUrlPtr
	bitratePtr
	outputFileNamePtr
	inputFileNamePtr
	heightPtr
	widthPtr
	frameRatePtr
	httpAddressPtr
	quantizationPtr
	timeoutPtr
	intraRefreshPeriodPtr
	verticalFlipPtr
	horizontalFlipPtr
	logPathPtr

	noOfConfigFlags
)

// Other misc consts
const (
	netSendRetryTime   = 5 * time.Second
	defaultRunDuration = 24 * time.Hour
	revidStopTime      = 5 * time.Second
	defaultLogPath     = "/var/log"
)

// canProfile is set to false with revid-cli is built with "-tags profile".
var canProfile = true

// The logger that will be used throughout
var logger *smartlogger.Logger

// Globals
var (
	rv          *revid.Revid
	config      revid.Config
	configFlags [](*string) = make([](*string), noOfConfigFlags)
)

func main() {
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to `file`")

	flagNames := [noOfConfigFlags]struct{ name, description string }{
		{"Input", "The input type: Raspivid, File"},
		{"InputCodec", "The codec of the input: H264, Mjpeg"},
		{"Output", "The output type: Http, Rtmp, File"},
		{"RtmpMethod", "The method used to send over rtmp: Ffmpeg, Librtmp"},
		// NOTE: we add rtp here when we have this functionality
		{"Packetization", "The method of data packetisation: Flv, Mpegts, None"},
		{"QuantizationMode", "Whether quantization if on or off (variable bitrate): On, Off"},
		{"Verbosity", "Verbosity: Info, Warning, Error, Fatal"},
		{"FramesPerClip", "Number of frames per clip sent"},
		{"RtmpUrl", "Url of rtmp endpoint"},
		{"Bitrate", "Bitrate of recorded video"},
		{"OutputFileName", "The directory of the output file"},
		{"InputFileName", "The directory of the input file"},
		{"Height", "Height in pixels"},
		{"Width", "Width in pixels"},
		{"FrameRate", "Frame rate of captured video"},
		{"HttpAddress", "Destination address of http posts"},
		{"Quantization", "Desired quantization value: 0-40"},
		{"Timeout", "Http timeout in seconds"},
		{"IntraRefreshPeriod", "The IntraRefreshPeriod i.e. how many keyframes we send"},
		{"VerticalFlip", "Flip video vertically: Yes, No"},
		{"HorizontalFlip", "Flip video horizontally: Yes, No"},
		{"LogPath", "Path for logging files (default is ../)"},
	}

	// Create the configFlags based on the flagNames struct array
	for i, f := range &flagNames {
		configFlags[i] = flag.String(f.name, "", f.description)
	}

	// Do we want a netsender session
	useNetsender := flag.Bool("NetSender", false, "Are we checking vars through netsender?")
	// User might also want to define how long revid runs for
	runDurationPtr := flag.Duration("runDuration", defaultRunDuration, "How long do you want revid to run for?")

	flag.Parse()

	handleFlags()

	err = run()
	if err != nil {
		logger.Log(Error, err.Error()) // TODO(kortschak): Make this "Fatal" when that exists.
		os.Exit(1)
	}
}

// initialize then run the main NetSender client
func run() error {
	// initialize NetSender and use NetSender's logger
	//config.Logger = netsender.Logger()
	logger.Log(smartlogger.Info, "Running in NetSender mode")

	var ns netsender.Sender
	err := ns.Init(nil, nil, nil, nil)
	if err != nil {
		return err
	}
	vars, _ := ns.Vars()
	vs := ns.VarSum()
	paused := false
	if vars["mode"] == "Paused" {
		paused = true
	}
	if !paused {
		err = updateRevid(&ns, vars, false)
		if err != nil {
			return err
		}
	}

	for {
		if err := send(&ns); err != nil {
			logger.Log(smartlogger.Warning, err.Error())
			time.Sleep(netSendRetryTime)
			continue
		}

		if vs != ns.VarSum() {
			// vars changed
			vars, err := ns.Vars()
			if err != nil {
				logger.Log(smartlogger.Warning, err.Error())
				time.Sleep(netSendRetryTime)
				continue
			}
			vs = ns.VarSum()
			if vars["mode"] == "Paused" {
				if !paused {
					logger.Log(Info, "Pausing revid")
					stopRevid()
					paused = true
				}
			} else {
				err = updateRevid(&ns, vars, !paused)
				if err != nil {
					return err
				}
				if paused {
					paused = false
				}
			}
		}
		sleepTime, _ := strconv.Atoi(ns.Param("mp"))
		time.Sleep(time.Duration(sleepTime) * time.Second)
	}
}

// handle flags interprets command line flags and sets revid config etc accordingly
func handleFlags() {
	logPath := defaultLogPath
	if *configFlags[logPathPtr] != "" {
		logPath = *configFlags[logPathPtr]
	}

	logger = smartlogger.New(defaultVerbosity, logPath)

	switch *configFlags[logVerbosityPtr] {
	case "Debug":
		config.Verbosity = smartlogger.Debug
	case "Info":
		config.Verbosity = smartlogger.Info
	case "Warning":
		config.Verbosity = smartlogger.Warning
	case "Error":
		config.Verbosity = smartlogger.Error
	case "Fatal":
		config.Verbosity = smartlogger.Fatal
	case "":
		config.Verbosity = defaultLogVerbosity
	default:
		logger.Log(smartlogger.Warning, "Bad input argument for LogVerbosity!")
	}

	if *cpuprofile != "" {
		if canProfile {
			f, err := os.Create(*cpuprofile)
			if err != nil {
				logger.Log(smartlogger.Fatal, "could not create CPU profile: ", err)
			}
			if err := pprof.StartCPUProfile(f); err != nil {
				logger.Log(smartlogger.Fatal, "could not start CPU profile: ", err)
			}
			defer pprof.StopCPUProfile()
		} else {
			logger.Log(smartlogger.Warning, "ignoring cpuprofile flag - http/pprof built in.")
		}
	}

	switch *configFlags[inputPtr] {
	case "Raspivid":
		config.Input = revid.Raspivid
	case "File":
		config.Input = revid.File
	case "":
	default:
		logger.Log(smartlogger.Error, "Bad input argument")
	}

	switch *configFlags[inputCodecPtr] {
	case "H264Codec":
		config.InputCodec = revid.H264Codec
	case "":
	default:
		logger.Log(smartlogger.Error, "Bad input codec argument!")
	}

	switch *configFlags[outputPtr] {
	case "File":
		config.Output = revid.File
	case "Http":
		config.Output = revid.Http
	case "Rtmp":
		config.Output = revid.Rtmp
	case "FfmpegRtmp":
		config.Output = revid.FfmpegRtmp
	case "":
	default:
		logger.Log(smartlogger.Error, "Bad output argument!")
	}

	switch *configFlags[rtmpMethodPtr] {
	case "Ffmpeg":
		config.RtmpMethod = revid.Ffmpeg
	case "LibRtmp":
		config.RtmpMethod = revid.LibRtmp
	case "":
	default:
		logger.Log(smartlogger.Error, "Bad rtmp method argument!")
	}

	switch *configFlags[packetizationPtr] {
	case "None":
		config.Packetization = revid.None
	case "Mpegts":
		config.Packetization = revid.Mpegts
	case "Flv":
		config.Packetization = revid.Flv
	case "":
	default:
		logger.Log(smartlogger.Error, "Bad packetization argument!")
	}

	switch *configFlags[quantizationModePtr] {
	case "QuantizationOn":
		config.QuantizationMode = revid.QuantizationOn
	case "QuantizationOff":
		config.QuantizationMode = revid.QuantizationOff
	case "":
	default:
		logger.Log(smartlogger.Error, "Bad quantization mode argument!")
	}

	switch *configFlags[verbosityPtr] {
	case "No":
		config.Verbosity = revid.No
	case "Yes":
		config.Verbosity = revid.Yes
	case "":
	default:
		logger.Log(smartlogger.Error, "Bad verbosity argument!")
	}

	switch *configFlags[horizontalFlipPtr] {
	case "No":
		config.HorizontalFlip = revid.No
	case "Yes":
		config.HorizontalFlip = revid.Yes
	case "":
		config.HorizontalFlip = revid.No
	default:
		logger.Log(smartlogger.Error, "Bad horizontal flip option!")
	}

	switch *configFlags[verticalFlipPtr] {
	case "No":
		config.VerticalFlip = revid.No
	case "Yes":
		config.VerticalFlip = revid.Yes
	case "":
		config.VerticalFlip = revid.No
	default:
		logger.Log(smartlogger.Error, "Bad vertical flip option!")
	}

	fpc, err := strconv.Atoi(*configFlags[framesPerClipPtr])
	if err == nil && fpc > 0 {
		config.FramesPerClip = fpc
	}
	config.RtmpUrl = *configFlags[rtmpUrlPtr]
	config.Bitrate = *configFlags[bitratePtr]
	config.OutputFileName = *configFlags[outputFileNamePtr]
	config.InputFileName = *configFlags[inputFileNamePtr]
	config.Height = *configFlags[heightPtr]
	config.Width = *configFlags[widthPtr]
	config.FrameRate = *configFlags[frameRatePtr]
	config.HttpAddress = *configFlags[httpAddressPtr]
	config.Quantization = *configFlags[quantizationPtr]
	config.Timeout = *configFlags[timeoutPtr]
	config.IntraRefreshPeriod = *configFlags[intraRefreshPeriodPtr]

	if !*useNetsender {
		// run revid for the specified duration
		startRevid(nil)
		time.Sleep(*runDurationPtr)
		stopRevid()
		return
	}
}

// send implements our main NetSender client and handles NetReceiver configuration
// (as distinct from httpSender which just sends video data).
func send(ns *netsender.Sender) error {
	// populate input values, if any
	inputs := netsender.MakePins(ns.Param("ip"), "X")
	if rv != nil {
		for i, pin := range inputs {
			if pin.Name == "X23" {
				inputs[i].Value = rv.Bitrate()
			}
		}
	}

	_, reconfig, err := ns.Send(netsender.RequestPoll, inputs)
	if err != nil {
		return err
	}
	if reconfig {
		return ns.Config()
	}
	return nil
}

// wrappers for stopping and starting revid
func startRevid(ns *netsender.Sender) (err error) {
	rv, err = revid.New(config, ns)
	if err != nil {
		return err
	}
	rv.Start()
	return nil
}

func stopRevid() {
	rv.Stop()

	// FIXME(kortschak): Is this waiting on completion of work?
	// Use a wait group and Wait method if it is.
	time.Sleep(revidStopTime)
}

func updateRevid(ns *netsender.Sender, vars map[string]string, stop bool) error {
	if stop {
		stopRevid()
	}

	//look through the vars and update revid where needed
	for key, value := range vars {
		switch key {
		case "Output":
			switch value {
			case "File":
				config.Output = revid.File
			case "Http":
				config.Output = revid.Http
			case "Rtmp":
				config.Output = revid.Rtmp
			case "FfmpegRtmp":
				config.Output = revid.FfmpegRtmp
			default:
				logger.Log(smartlogger.Warning, "Invalid Output param: ", value)
				continue
			}
		case "FramesPerClip":
			fpc, err := strconv.Atoi(value)
			if fpc > 0 && err == nil {
				config.FramesPerClip = fpc
			} else {
				logger.Log(smartlogger.Warning, "Invalid FramesPerClip param: ", value)
			}
		case "RtmpUrl":
			config.RtmpUrl = value
		case "Bitrate":
			asInt, err := strconv.Atoi(value)
			if asInt > 0 && err == nil {
				config.Bitrate = value
			} else {
				logger.Log(smartlogger.Warning, "Invalid Bitrate param: ", value)
			}
		case "OutputFileName":
			config.OutputFileName = value
		case "InputFileName":
			config.InputFileName = value
		case "Height":
			asInt, err := strconv.Atoi(value)
			if asInt > 0 && err == nil {
				config.Height = value
			} else {
				logger.Log(smartlogger.Warning, "Invalid Height param: ", value)
			}
		case "Width":
			asInt, err := strconv.Atoi(value)
			if asInt > 0 && err == nil {
				config.Width = value
			} else {
				smartlogger.Log(smartlogger.Warning, "Invalid Width param: ", value)
			}
		case "FrameRate":
			asInt, err := strconv.Atoi(value)
			if asInt > 0 && err == nil {
				config.FrameRate = value
			} else {
				smartlogger.Log(smartlogger.Warning, "Invalid FrameRate param: ", value)
			}
		case "HttpAddress":
			config.HttpAddress = value
		case "Quantization":
			asInt, err := strconv.Atoi(value)
			if asInt > 0 && err == nil {
				config.Quantization = value
			} else {
				smartlogger.Log(smartlogger.Warning, "Invalid Quantization param: ", value)
			}
		case "Timeout":
			asInt, err := strconv.Atoi(value)
			if asInt > 0 && err == nil {
				config.Timeout = value
			}
		case "IntraRefreshPeriod":
			asInt, err := strconv.Atoi(value)
			if asInt > 0 && err == nil {
				config.IntraRefreshPeriod = value
			}
		case "HorizontalFlip":
			switch value {
			case "Yes":
				config.HorizontalFlip = revid.Yes
			case "No":
				config.HorizontalFlip = revid.No
			default:
				smartlogger.Log(smartlogger.Warning, "Invalid HorizontalFlip param: ", value)
			}
		case "VerticalFlip":
			switch value {
			case "Yes":
				config.VerticalFlip = revid.Yes
			case "No":
				config.VerticalFlip = revid.No
			default:
				smartlogger.Log(smartlogger.Warning, "Invalid VerticalFlip param: ", value)
			}
		default:
			if unicode.IsUpper(rune(key[0])) {
				smarlogger.Log(smartlogger.Warning, "Unexpected param: ", key)
			} // else system params are lower case
		}
	}

	return startRevid(ns)
}