package goconv

/*
#cgo CFLAGS: -I./ -I./inc -I/usr/local/cuda/include
#cgo CXXFLAGS: -I./ -I./inc -I/usr/local/cuda/include -std=c++11
#cgo LDFLAGS: -L/usr/local/cuda/lib64 -lnppi -lcudart -ldl
#include <stdlib.h>
#include "conv.h"
*/
import "C"
import (
	"unsafe"

	"basic.com/valib/godraw.git"
	"basic.com/valib/gogpu.git"
	"github.com/disintegration/imaging"
)

const (
	need     = 200
	reserved = 512
)

func gpuIndex(lastIndex int) int {
	indices := gogpu.RankGPU()
	if len(indices) == 0 {
		return -1
	}

	for _, v := range indices {
		if v != lastIndex {
			if gogpu.SatisfyGPU(v, need, need/2) {
				return v
			}
		}
	}

	if gogpu.SatisfyGPU(lastIndex, need, reserved) {
		return lastIndex
	}
	return -1
}

type convertor struct {
	width   int
	height  int
	rWidth  int
	rHeight int
	conv    C.convHandle
}

var convts []*convertor

func find(w, h, rw, rh int) *convertor {
	for _, v := range convts {
		if v.width == w && v.height == h && v.rWidth == rw && v.rHeight == rh {
			return v
		}
	}
	gpu := gpuIndex(0)
	if gpu < 0 {
		return nil
	}
	cw := C.conv_create(C.int(w), C.int(h), C.int(rw), C.int(rh), C.int(gpu))
	if cw == nil {
		return nil
	}
	c := &convertor{w, h, rw, rh, cw}
	convts = append(convts, c)
	return c
}

// YUV2BGR yuv->bgr
func YUV2BGR(yuv []byte, w, h int) []byte {

	cw := find(w, h, -1, -1)
	if cw == nil {
		return yuv2bgr(yuv, w, h)
	}
	var bgr *C.uchar
	var bgrLen C.int
	ret := C.yuv2bgr(cw.conv, unsafe.Pointer(&yuv[0]), &bgr, &bgrLen)
	if ret != 0 {
		return nil
	}
	const maxLen = 0x7fffffff
	goBGRLen := int(bgrLen)
	if goBGRLen > 0 {
		return (*[maxLen]byte)(unsafe.Pointer(bgr))[:goBGRLen:goBGRLen]
	}
	return nil
}

// YUV2ResizedBGR yuv -> resized bgr
func YUV2ResizedBGR(yuv []byte, w, h, rw, rh int) []byte {

	cw := find(w, h, rw, rh)
	if cw == nil {
		bgr := yuv2bgr(yuv, w, h)
		return bgresize(bgr, w, h, rw, rh)
	}
	var bgr *C.uchar
	var bgrLen C.int
	ret := C.yuv2resizedbgr(cw.conv, unsafe.Pointer(&yuv[0]), &bgr, &bgrLen)
	if ret != 0 {
		return nil
	}
	const maxLen = 0x7fffffff
	goBGRLen := int(bgrLen)
	if goBGRLen > 0 {
		return (*[maxLen]byte)(unsafe.Pointer(bgr))[:goBGRLen:goBGRLen]
	}
	return nil

}

// ResizeBGR resize
func ResizeBGR(bgrO []byte, w, h, rw, rh int) []byte {
	if (rw < 0 && rh < 0) || (rw > w && rh > h) {
		return bgrO
	}

	cw := find(w, h, rw, rh)
	if cw == nil {
		return bgresize(bgrO, w, h, rw, rh)
	}

	var bgr *C.uchar
	var bgrLen C.int
	ret := C.resizebgr(cw.conv, unsafe.Pointer(&bgrO[0]), &bgr, &bgrLen)
	if ret != 0 {
		return nil
	}
	const maxLen = 0x7fffffff
	goBGRLen := int(bgrLen)
	if goBGRLen > 0 {
		return (*[maxLen]byte)(unsafe.Pointer(bgr))[:goBGRLen:goBGRLen]
	}
	return nil
}

// ResizeYUV yuv
func ResizeYUV(yuv []byte, w, h, rw, rh int) []byte {
	if (rw < 0 && rh < 0) || (rw > w && rh > h) {
		return yuv
	}

	cw := find(w, h, rw, rh)
	if cw == nil {
		return yuv
	}

	var resized *C.uchar
	var resizedLen C.int
	ret := C.resizeyuv(cw.conv, unsafe.Pointer(&yuv[0]), &resized, &resizedLen)
	if ret != 0 {
		return nil
	}

	const maxLen = 0x7fffffff
	goResizedLen := int(resizedLen)
	if goResizedLen > 0 {
		return (*[maxLen]byte)(unsafe.Pointer(resized))[:goResizedLen:goResizedLen]
	}
	return nil
}

// YUV2BGRandResize conv and resize
func YUV2BGRandResize(yuv []byte, w, h, rw, rh int) ([]byte, []byte) {
	cw := find(w, h, rw, rh)
	if cw == nil {
		origin := yuv2bgr(yuv, w, h)
		resized := bgresize(origin, w, h, rw, rh)
		return origin, resized
	}

	var bgr *C.uchar
	var bgrLen C.int
	var scale *C.uchar
	var scaleLen C.int

	ret := C.yuv2bgrandresize(cw.conv, unsafe.Pointer(&yuv[0]), &bgr, &bgrLen, &scale, &scaleLen)

	if ret != 0 {
		return nil, nil
	}
	var out, resized []byte

	const maxLen = 0x7fffffff
	goBGRLen, goScaleLen := int(bgrLen), int(scaleLen)
	if goBGRLen > 0 {
		out = (*[maxLen]byte)(unsafe.Pointer(bgr))[:goBGRLen:goBGRLen]
	}
	if goScaleLen > 0 {
		resized = (*[maxLen]byte)(unsafe.Pointer(scale))[:goScaleLen:goScaleLen]
	}
	return out, resized

}

// Free free
func Free() {
	for _, v := range convts {
		if v.conv != nil {
			C.conv_destroy(v.conv)
		}
	}
}

func yuv2bgr(yuv []byte, w, h int) []byte {

	data := make([]byte, 0, w*h*3)
	start := w * h
	for i := 0; i < h; i++ {
		for j := 0; j < w; j++ {

			index := i/2*w + j - (j & 0x01)

			y := int32(yuv[j+i*w])
			v := int32(yuv[start+index])
			u := int32(yuv[start+index+1])

			r := y + (140*(v-128))/100
			g := y - (34*(u-128)+71*(v-128))/100
			b := y + (177*(u-128))/100

			if r > 255 {
				r = 255
			}
			if r < 0 {
				r = 0
			}
			if g > 255 {
				g = 255
			}
			if g < 0 {
				g = 0
			}
			if b > 255 {
				b = 255
			}
			if b < 0 {
				b = 0
			}
			data = append(data, byte(r), byte(g), byte(b))
		}
	}
	return data
}

func bgresize(bgr []byte, w, h, rw, rh int) []byte {
	img, err := godraw.ToImage(bgr, w, h)
	if err != nil {
		return nil
	}
	dstImg := imaging.Resize(img, rw, rh, imaging.NearestNeighbor)
	return godraw.Image2BGR(dstImg)
}