package sdk

import (
	"analysis/logo"
	"analysis/work"
	"context"
	"fmt"
	"plugin"

	"basic.com/libgowrapper/sdkstruct.git"
	"basic.com/pubsub/protomsg.git"
	"basic.com/valib/gogpu.git"
	"github.com/gogo/protobuf/proto"
)

// YoloDetect yolo detect
type YoloDetect struct {
	iGPU    int
	cfg     string
	weights string
	name    string

	list *LockList

	handle interface{}
	fnInit func(string, string, string, int, func(...interface{})) interface{}
	fnFree func(interface{})
	fnRun  func(interface{}, string, []byte, int, int, int, float32, int) ([]sdkstruct.CObjTrackInfo, []sdkstruct.CObjTrackInfo)
}

// NewYDetectWithTrack with track
func NewYDetectWithTrack(gi int, cfg, weights, name string) *YoloDetect {
	soFile := "libyolo.so"

	plug, err := plugin.Open(soFile)
	if err != nil {
		logo.Errorln("Open: ", soFile, " error: ", err)
		return nil
	}

	fnInit, _ := LoadFunc(plug, soFile, "NewSDK")
	fnFree, _ := LoadFunc(plug, soFile, "Free")
	fnRun, _ := LoadFunc(plug, soFile, "Run")

	return &YoloDetect{
		iGPU:    gi,
		cfg:     cfg,
		weights: weights,
		name:    name,

		list: NewLockList(6),

		handle: nil,
		fnInit: fnInit.(func(string, string, string, int, func(...interface{})) interface{}),
		fnFree: fnFree.(func(interface{})),
		fnRun:  fnRun.(func(interface{}, string, []byte, int, int, int, float32, int) ([]sdkstruct.CObjTrackInfo, []sdkstruct.CObjTrackInfo)),
	}
}

// Init impl interface
func (y *YoloDetect) Init() bool {

	gpu := y.iGPU
	if gpu == -1 {
		gpu = gogpu.ValidGPU(2048)
	}
	h := y.fnInit(y.cfg, y.weights, y.name, gpu, logo.Infoln)
	logo.Infoln("YOLO USE GPU: ", gpu)

	if h == nil {
		logo.Errorln("CREATE YOLO DETECTOR ERROR")
		return false
	}

	y.handle = h
	return true
}

// Run impl interface
func (y *YoloDetect) Run(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {
	FlowSimple(ctx, in, out, typ, y.list.Push, y.list.Pop, y.track, func() { y.fnFree(y.handle) })
}

func (y *YoloDetect) track(rMsg work.MsgRS, out chan<- work.MsgRS, typ string) {

	/////////////////////////////////////////////
	i := unpackImage(rMsg, typ)
	if i == nil || i.Data == nil || i.Width <= 0 || i.Height <= 0 {
		logo.Errorln("yolo image error: ", i)
		ejectResult(nil, rMsg, out)

		return
	}

	imgW, imgH := int(i.Width), int(i.Height)

	whole, recent := y.fnRun(y.handle, rMsg.Msg.Cid, i.Data, imgW, imgH, 3, 0.4, 0)
	if len(recent) > 0 {
	}

	infos := convert2ProtoYoloTrack(whole, 1.0, 1.0)
	p := protomsg.ParamYoloObj{Infos: infos}

	data, err := proto.Marshal(&p)
	if err != nil {
		fmt.Println("ydetect track marshal proto yolo obj error", err)
		data = nil
	}

	ejectResult(data, rMsg, out)

	var id, name string
	if rMsg.Msg.Tasklab != nil {
		id, name = rMsg.Msg.Tasklab.Taskid, rMsg.Msg.Tasklab.Taskname
	}

	logo.Infoln("CAMERAID: ", rMsg.Msg.Cid, " TASKID: ", id, " TASKNAME: ", name, " DETECT YOLO COUNT: ", len(whole))
}

func convert2ProtoYoloTrack(obj []sdkstruct.CObjTrackInfo, fx, fy float64) []*protomsg.ObjInfo {
	ret := []*protomsg.ObjInfo{}

	for _, v := range obj {
		if fx < 1.0 || fy < 1.0 {
			v.ObjInfo.RcObj.Left = (int32)((float64)(v.ObjInfo.RcObj.Left) / fx)
			v.ObjInfo.RcObj.Right = (int32)((float64)(v.ObjInfo.RcObj.Right) / fx)
			v.ObjInfo.RcObj.Top = (int32)((float64)(v.ObjInfo.RcObj.Top) / fy)
			v.ObjInfo.RcObj.Bottom = (int32)((float64)(v.ObjInfo.RcObj.Bottom) / fy)
		}

		rect := protomsg.Rect{
			Left:   v.ObjInfo.RcObj.Left,
			Right:  v.ObjInfo.RcObj.Right,
			Top:    v.ObjInfo.RcObj.Top,
			Bottom: v.ObjInfo.RcObj.Bottom,
		}
		obj := protomsg.ObjInfo{
			RcObj: &rect,
			Typ:   v.ObjInfo.Typ,
			Prob:  v.ObjInfo.Prob,
			ObjID: v.ID,
		}

		ret = append(ret, &obj)
	}
	return ret
}