package main

import (
	"bufio"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"
	"unsafe"

	"basic.com/libgowrapper/sdkhelper.git"
	"basic.com/valib/gogpu.git"

	"basic.com/pubsub/protomsg.git"
	"basic.com/valib/godraw.git"

	"github.com/gogo/protobuf/proto"

	"nanomsg.org/go-mangos"
	"nanomsg.org/go-mangos/protocol/rep"
	"nanomsg.org/go-mangos/transport/ipc"
	"nanomsg.org/go-mangos/transport/tcp"
)

type reid struct {
	handle   *ReID
	fnLogger func(...interface{})

	gpu int
	ipc string
}

// Create Reid
func Create(config string, typ, id string, gpu int, shm bool, ipc2Rule string, ruleMaxSize int, fn func(...interface{}), reserved map[string]interface{}) interface{} {

	cfg, err := sdkhelper.ReadConfig(config)
	if err != nil {
		fn("Reid SDK Create Error When Read Config: ", err)
		return nil
	}

	sModel, sGPU, sIPC :=
		"reid-model",
		"gpu-memory",
		"ipc-url"

	params := []string{sModel, sGPU, sIPC}

	for _, v := range params {
		if _, ok := cfg.Param[v]; !ok {
			fn("Reid SDK Create Error Because of Param Not Found: ", v)
			return nil
		}
	}

	gpuM := sdkhelper.Atoi(cfg.Param[sGPU])

	rGPU := gpu

	if rGPU == -1 {
		rGPU = gogpu.ValidGPU(gpuM + 512)
	}

	handle := NewSDK(rGPU, cfg.Param[sModel])
	if handle == nil {
		fn("Reid SDK Create Error When New SDK")
		return nil
	}

	return &reid{
		handle:   handle,
		fnLogger: fn,

		gpu: rGPU,
		ipc: cfg.Param[sIPC],
	}
}

// Run run
func Run(ctx context.Context, i interface{}) {
	s := i.(*reid)

	var sock mangos.Socket
	var err error
	var msg []byte

	for {
		if sock, err = rep.NewSocket(); err != nil {
			s.fnLogger("can't get new rep socket: ", err)
			time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
	}

	sock.AddTransport(ipc.NewTransport())
	sock.AddTransport(tcp.NewTransport())

	for {
		if err = sock.Listen(s.ipc); err != nil {
			suf := "ipc://"
			p := strings.Index(s.ipc, suf)
			if p >= 0 {
				file := s.ipc[p+len(string(suf)):]
				os.Remove(file)
                s.fnLogger("remove:", file)
			}
			s.fnLogger("can't listen on rep socket: ", err)
			time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
	}

	for {
		select {
		case <-ctx.Done():
			sock.Close()
			return
		default:

			if msg, err = sock.Recv(); err != nil {
				s.fnLogger("REID~~~~~~Recv From HumanTrack error: ", err)
				continue
			}

			i := &protomsg.Image{}
			err := proto.Unmarshal(msg, i)
			if err != nil {
				s.fnLogger("REID~~~~~~protobuf decode CameraImage error: ", err)
				continue
			}
			if i.Data == nil {
				s.fnLogger("REID~~~~~~protomsg.Image data null")
				continue
			}
			s.fnLogger("REID~~~~~~Recv Image:", len(i.Data))

			/////////////////////////////////
			if jpg, err := godraw.ToJpeg(i.Data, int(i.Width), int(i.Height), nil); err == nil {
				ioutil.WriteFile("./pic.jpg", jpg, 0644)
			}
			/////////////////////////////////
			feat := s.handle.Extract2(unsafe.Pointer(&i.Data[0]), int(i.Width), int(i.Height), 3)
			if feat == nil {
				// feat = make([]float32, 1)
			} else {
				for k := 0; k < 3; k++ {
					s.fnLogger("REID~~~~~~extractor---human_feats------%f", feat[k+2000])
				}
				s.fnLogger("REID~~~~~~Run Reid Use GPU: ", s.gpu)

				/////////////////////////////////
				if f, err := os.Create("./reid-feat.txt"); err == nil {
					defer f.Close()

					w := bufio.NewWriter(f)
					for k, v := range feat {
						lineStr := fmt.Sprintf("%d->%f", k, v)
						fmt.Fprintln(w, lineStr)
					}
					w.Flush()
				}
				/////////////////////////////////

			}
			buf := float32SliceAsByteSlice(feat)
			ioutil.WriteFile("./reid-feat-byte.txt", buf, 0644)

			if err = sock.Send(buf); err != nil {
				s.fnLogger("REID~~~~~~Send HumanTrack error: ", err)
			}

		}

	}
}

func float32SliceAsByteSlice(src []float32) []byte {
	if len(src) == 0 {
		return nil
	}

	l := len(src) * 4
	ptr := unsafe.Pointer(&src[0])
	return (*[1 << 26]byte)((*[1 << 26]byte)(ptr))[:l:l]
}