package main

/*
#include <stdlib.h>
#include <stdio.h>

static int hw_pix_fmt;

static inline int get_hw_pix_fmt()
{
	return hw_pix_fmt;
}

static inline void set_hw_pix_fmt(int fmt)
{
	hw_pix_fmt = fmt;
}

struct AVCodecContext;

int get_hw_format(struct AVCodecContext *ctx, const int *pix_fmts)
{
	const int *p;

	fprintf(stderr, "get_hw_format called.\n");
	for (p = pix_fmts; *p != -1; p++) {
		if (*p == hw_pix_fmt)
			return *p;
	}

	fprintf(stderr, "Failed to get HW surface format.\n");
	return -1;
}
*/
import "C"

import (
	"fmt"
	"os"
	"syscall"
	"unsafe"

	"github.com/qrtc/ffmpeg-dev-go"
)

var (
	hwDeviceCtx *ffmpeg.AVBufferRef
	outputFile  *os.File
)

func hwDecoderInit(ctx *ffmpeg.AVCodecContext, hwType ffmpeg.AVHWDeviceType) (ret int32) {
	if ret := ffmpeg.AvHWDeviceCtxCreate(&hwDeviceCtx, hwType, "", nil, 0); ret < 0 {
		fmt.Fprintf(os.Stderr, "Failed to create specified HW device.\n")
		return ret
	}
	ctx.SetHwDeviceCtx(ffmpeg.AvBufferRef(hwDeviceCtx))
	return ret
}

func decodeWrite(avctx *ffmpeg.AVCodecContext, packet *ffmpeg.AVPacket) (ret int32) {
	var frame, swFrame, tmpFrame *ffmpeg.AVFrame
	var buffer *uint8
	var size int32

	if ret = ffmpeg.AvCodecSendPacket(avctx, packet); ret < 0 {
		fmt.Fprintf(os.Stderr, "Error during decoding\n")
		return ret
	}

	for {
		if frame = ffmpeg.AvFrameAlloc(); frame == nil {
			fmt.Fprintf(os.Stderr, "Can not alloc frame\n")
			goto fail
		}
		if swFrame = ffmpeg.AvFrameAlloc(); swFrame == nil {
			fmt.Fprintf(os.Stderr, "Can not alloc sw frame\n")
			goto fail
		}

		ret = ffmpeg.AvCodecReceiveFrame(avctx, frame)
		if ret == ffmpeg.AVERROR(syscall.EAGAIN) || ret == ffmpeg.AVERROR_EOF {
			ffmpeg.AvFrameFree(&frame)
			ffmpeg.AvFrameFree(&swFrame)
			return 0
		} else if ret < 0 {
			fmt.Fprintf(os.Stderr, "Error while decoding\n")
			goto fail
		}

		if frame.GetFormat() == (int32)(C.get_hw_pix_fmt()) {
			// retrieve data from GPU to CPU
			if ret = ffmpeg.AvHWFrameTransferData(swFrame, frame, 0); ret < 0 {
				fmt.Fprintf(os.Stderr, "Error transferring the data to system memory\n")
				goto fail
			}
			tmpFrame = swFrame
		} else {
			tmpFrame = frame
		}

		size = ffmpeg.AvImageGetBufferSize(tmpFrame.GetFormat(), tmpFrame.GetWidth(),
			tmpFrame.GetHeight(), 1)
		buffer = (*uint8)(ffmpeg.AvMalloc(size))
		if buffer == nil {
			fmt.Fprintf(os.Stderr, "Can not alloc buffer\n")
			ret = ffmpeg.AVERROR(syscall.ENOMEM)
			goto fail
		}

		ret = ffmpeg.AvImageCopyToBuffer(buffer, size, tmpFrame.GetData(),
			tmpFrame.GetLinesize(), tmpFrame.GetFormat(),
			tmpFrame.GetWidth(), tmpFrame.GetHeight(), 1)
		if ret < 0 {
			fmt.Fprintf(os.Stderr, "Can not copy image to buffer\n")
			goto fail
		}
		if _, err := outputFile.Write(unsafe.Slice(buffer, size)); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to dump raw data.\n")
			goto fail
		}
	fail:
		ffmpeg.AvFrameFree(&frame)
		ffmpeg.AvFrameFree(&swFrame)
		ffmpeg.AvFreep(&buffer)
		if ret < 0 {
			return ret
		}
	}
}

func main() {
	var inputCtx *ffmpeg.AVFormatContext
	var videoStream, ret int32
	var video *ffmpeg.AVStream
	var decoderCtx *ffmpeg.AVCodecContext
	var decoder *ffmpeg.AVCodec
	var packet ffmpeg.AVPacket
	var hwType ffmpeg.AVHWDeviceType

	if len(os.Args) < 4 {
		fmt.Fprintf(os.Stderr, "Usage: %s <device type> <input file> <output file>\n", os.Args[0])
		os.Exit(1)
	}

	hwType = ffmpeg.AvHWDeviceFindTypeByName(os.Args[1])
	if hwType == ffmpeg.AV_HWDEVICE_TYPE_NONE {
		fmt.Fprintf(os.Stderr, "Device type %s is not supported.\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Available device types:")
		for hwType = ffmpeg.AvHWDeviceIterateTypes(hwType); hwType != ffmpeg.AV_HWDEVICE_TYPE_NONE; {
			fmt.Fprintf(os.Stderr, " %s", ffmpeg.AvHWDeviceGetTypeName(hwType))
			hwType = ffmpeg.AvHWDeviceIterateTypes(hwType)
		}
		fmt.Fprintf(os.Stderr, "\n")
		os.Exit(1)
	}

	// open the input file
	if ret = ffmpeg.AvFormatOpenInput(&inputCtx, os.Args[2], nil, nil); ret != 0 {
		fmt.Fprintf(os.Stderr, "Cannot open input file '%s'\n", os.Args[2])
		os.Exit(1)
	}

	if ret = ffmpeg.AvFormatFindStreamInfo(inputCtx, nil); ret < 0 {
		fmt.Fprintf(os.Stderr, "Cannot find input stream information.\n")
		os.Exit(1)
	}

	// find the video stream information
	ret = ffmpeg.AvFindBestStream(inputCtx, ffmpeg.AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0)
	if ret < 0 {
		fmt.Fprintf(os.Stderr, "Cannot find a video stream in the input file\n")
		os.Exit(1)
	}
	videoStream = ret

	for i := 0; ; i++ {
		config := ffmpeg.AvCodecGetHwConfig(decoder, i)
		if config == nil {
			fmt.Fprintf(os.Stderr, "Decoder %s does not support device type %s.\n",
				decoder.GetName(), ffmpeg.AvHWDeviceGetTypeName(hwType))
			os.Exit(1)
		}
		if (config.GetMethods()&ffmpeg.AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) != 0 &&
			config.GetDeviceType() == hwType {
			fmt.Fprintf(os.Stdout, "Support set_hw_pix_fmt. \n")
			C.set_hw_pix_fmt((C.int)(config.GetPixFmt()))
			break
		}
	}

	if decoderCtx = ffmpeg.AvCodecAllocContext3(decoder); decoderCtx == nil {
		os.Exit(int(ffmpeg.AVERROR(syscall.ENOMEM)))
	}

	video = inputCtx.GetStreamsIdx(int(videoStream))
	if ret = ffmpeg.AvCodecParametersToContext(decoderCtx, video.GetCodecpar()); ret < 0 {
		os.Exit(1)
	}

	decoderCtx.SetGetFormat((ffmpeg.AVCodecContextGetFormatFunc)(C.get_hw_format))

	if ret = hwDecoderInit(decoderCtx, hwType); ret < 0 {
		os.Exit(1)
	}

	if ret = ffmpeg.AvCodecOpen2(decoderCtx, decoder, nil); ret < 0 {
		fmt.Fprintf(os.Stderr, "Failed to open codec for stream #%d\n", videoStream)
		os.Exit(1)
	}

	// open the file to dump raw data
	outputFile, _ = os.OpenFile(os.Args[3], os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)

	// actual decoding and dump the raw data
	for ret >= 0 {
		if ret = ffmpeg.AvReadFrame(inputCtx, &packet); ret < 0 {
			break
		}
		if videoStream == packet.GetStreamIndex() {
			ret = decodeWrite(decoderCtx, &packet)
		}

		ffmpeg.AvPacketUnref(&packet)
	}

	// flush the decoder
	packet.SetData(nil)
	packet.SetSize(0)
	ret = decodeWrite(decoderCtx, &packet)
	ffmpeg.AvPacketUnref(&packet)

	if outputFile != nil {
		outputFile.Close()
	}
	ffmpeg.AvCodecFreeContext(&decoderCtx)
	ffmpeg.AvFormatCloseInput(&inputCtx)
	ffmpeg.AvBufferUnref(&hwDeviceCtx)

	if ret != 0 {
		fmt.Fprintf(os.Stderr, "(error '%s')\n", ffmpeg.AvErr2str(ret))
	}
	os.Exit(int(ret))
}