package sdk

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

	"github.com/gogo/protobuf/proto"

	"basic.com/libgowrapper/sdkstruct.git"
	"basic.com/pubsub/protomsg.git"
	"basic.com/valib/gogpu.git"
)

// HumanTracker track
type HumanTracker struct {
	gpu       int
	batchSize int
	flag      int

	list *LockList

	handle       interface{}
	fnInit       func(int, int, int, func(...interface{})) interface{}
	fnFree       func(interface{})
	fnRun        func(interface{}, []byte, int, int, int) []sdkstruct.FgResult
	fnProcess    func(interface{}, []sdkstruct.SDKImage) ([]sdkstruct.FgResult, error)
	fnVer        func(interface{}) string
	fnSimilarity func([]float32, []float32) (float64, error)
}

// NewHumanTracker new
func NewHumanTracker(gpu, batchSize, flag int) *HumanTracker {
	soFile := "libhumantrack.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")
	fnProc, _ := LoadFunc(plug, soFile, "Process")
	fnVer, _ := LoadFunc(plug, soFile, "GetVersion")
	fnSim, _ := LoadFunc(plug, soFile, "CFFSimilarity")

	return &HumanTracker{
		gpu:       gpu,
		batchSize: batchSize,
		flag:      flag,

		list: NewLockList(6),

		handle:       nil,
		fnInit:       fnInit.(func(int, int, int, func(...interface{})) interface{}),
		fnFree:       fnFree.(func(interface{})),
		fnRun:        fnRun.(func(interface{}, []byte, int, int, int) []sdkstruct.FgResult),
		fnProcess:    fnProc.(func(interface{}, []sdkstruct.SDKImage) ([]sdkstruct.FgResult, error)),
		fnVer:        fnVer.(func(interface{}) string),
		fnSimilarity: fnSim.(func([]float32, []float32) (float64, error)),
	}
}

// Init impl
func (t *HumanTracker) Init() bool {
	if t.batchSize != 1 {
		logo.Errorln("ONLY SUPPORT BATCH SIZE = 1")
		return false
	}
	gpu := t.gpu
	if gpu == -1 {
		gpu = gogpu.ValidGPU(2048)
	}
	h := t.fnInit(gpu, t.batchSize, t.flag, logo.Infoln)
	logo.Infoln("HumanTrack USE GPU: ", gpu)

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

	t.handle = h
	return true
}

func (t *HumanTracker) 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 {
		ejectResult(nil, rMsg, out)
		return
	}

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

	// var images []sdkstruct.SDKImage
	// img := sdkstruct.SDKImage{
	// 	Data:    i.Data,
	// 	Width:   imgW,
	// 	Height:  imgH,
	// 	Channel: 3,
	// }
	// images = append(images, img)
	res := t.fnRun(t.handle, i.Data, imgW, imgH, 3)
	if res == nil {
		logo.Errorln("HUMAN TRACKER RUN ERROR")

		ejectResult(nil, rMsg, out)
		return
	}

	hr := convert2ProtoHumanTrackResult(res)
	result := protomsg.HumanTrackResult{Result: hr[0]}
	data, err := proto.Marshal(&result)
	if err != nil {
		logo.Errorln("HUMAN TRACKER MARSHAL PROTO 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(hr[0]))

}

// Run impl
func (t *HumanTracker) Run(ctx context.Context, in <-chan work.MsgRS, out chan<- work.MsgRS, typ string) {

	FlowSimple(ctx, in, out, typ, t.list.Push, t.list.Pop, t.track, func() { t.fnFree(t.handle) })
}

// message HumanTrack {
//     Rect rcHuman = 1;
// 	float confidence = 2;
// 	int32 x = 3;
// 	int32 y = 4;
// 	int32 id = 5;
// 	repeated float feature = 6;
// }

// message HumanTrackResult {
//     repeated HumanTrack result = 1;
// }

func convert2ProtoHumanTrackResult(obj []sdkstruct.FgResult) [][]*protomsg.HumanTrack {
	ret := [][]*protomsg.HumanTrack{}
	for _, v := range obj {
		res := []*protomsg.HumanTrack{}
		for i := 0; i < int(v.FgNum); i++ {
			r := v.Fginfo[i]

			rect := protomsg.Rect{
				Left:   r.Left,
				Right:  r.Right,
				Top:    r.Top,
				Bottom: r.Bottom,
			}
			pr := &protomsg.HumanTrack{
				RcHuman:    &rect,
				Confidence: r.Confidence,
				X:          r.X,
				Y:          r.Y,
				Id:         r.ID,
				Feature:    r.Feature[:],
			}
			res = append(res, pr)
		}
		ret = append(ret, res)
	}
	return ret
}