package service

import (
	"errors"
	"fmt"
	"time"

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

// StreamMode stream mode
type StreamMode int

const (
	// Rtsp mode
	Rtsp = StreamMode(iota)
	// GB28181 mode
	GB28181
)

// PicExt file extention
type PicExt string

// PNGFileExt FileExt = ".png"
// JPEGFileExt FileExt = ".jpg"
// GIFFileExt FileExt = ".gif"

const (
	// PNGFileExt png
	PNGFileExt PicExt = PicExt(".png")
	// JPEGFileExt jpeg
	JPEGFileExt PicExt = PicExt(".jpg")
	// GIFFileExt gif
	GIFFileExt PicExt = PicExt(".gif")
	// InvalidFileExt invalid
	InvalidFileExt PicExt = PicExt("")
)

// capture pic
func captureGB28181(soFile string, url string, maxTry int) ([]byte, error) {

	var ret []byte
	var err error

	err = goffmpeg.InitFFmpeg(soFile)
	if err != nil {
		return nil, err
	}

	reterr := errors.New(fmt.Sprintf("gb28181 try %d times to capture image, is url correct?", maxTry))
	for i := 0; i < maxTry; i++ {
		ret = goffmpeg.GetGBJpg(url, maxTry)
		if len(ret) > 0 {
			reterr = nil
			break
		}
	}

	goffmpeg.FreeFFmpeg()

	return ret, reterr
}

// Capture pic
func Capture(soFile string, m StreamMode, url string, ext PicExt, w, h, maxTry int) ([]byte, error) {
	if m != GB28181 && m != Rtsp {
		return nil, errors.New("there is no this mode, try capture.Rtsp/capture.GB28181")
	}

	if m == GB28181 {
		return captureGB28181(soFile, url, maxTry)
	}

	if ext == InvalidFileExt {
		ext = JPEGFileExt
	}
	var ret []byte
	var err error

	err = goffmpeg.InitFFmpeg(soFile)
	if err != nil {
		return nil, err
	}

	ret, err = capt2Data(url, m, ext, w, h, maxTry)

	goffmpeg.FreeFFmpeg()

	return ret, err
}

func capt2Data(url string, m StreamMode, ext PicExt, w, h, maxTry int) ([]byte, error) {

	gb := false
	if m == GB28181 {
		gb = true
	}

	gf := goffmpeg.New(gb, false)
	defer gf.Free()

	gf.Run(url)
	gf.BuildDecoder()

	tryTime := 0
	// 保证有一个关键帧
	maxTry += 50
	if maxTry > 60 {
		maxTry = 60
	}

	var bgrData, jpgData []byte
	var err error
	for {

		data, ow, oh, _ := gf.GetYUV()
		if ow > 0 && oh > 0 {
			if ow*oh*3 != len(data) {
				// data 不是 bgr 格式，需要转换
				bgrData = yuv2bgr(data, ow, oh)
			} else {
				bgrData = data
			}

			if w > 0 && h > 0 {
				bgrData = bgresize(data, ow, oh, w, h)
				ow, oh = w, h
			}

			jpgData, err = godraw.ToJpeg(bgrData, ow, oh, nil)
			if err != nil {
				continue
			}
			break
		} else {
			tryTime++
			if tryTime > maxTry {
				break
			}
			time.Sleep(800 * time.Millisecond)
		}
	}

	if tryTime >= maxTry {
		err := fmt.Sprintf("try %d times to capture image, is url correct?\n", tryTime)
		return nil, errors.New(err)
	}

	return jpgData, nil
}

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)
}