package movielens

import (
	"context"
	"database/sql"

	"github.com/auxten/go-ctr/feature/ubcache"
	rcmd "github.com/auxten/go-ctr/recommend"
	"github.com/auxten/go-ctr/utils"
	log "github.com/sirupsen/logrus"
)

func (recSys *MovielensRec) GetUserBehavior(ctx context.Context, userId int,
	maxLen int64, maxPk int64, maxTs int64) (itemSeq []int, err error) {
	var (
		rows      *sql.Rows
		tableName string
	)

	// get stage value from ctx
	stage := ctx.Value(rcmd.StageKey).(rcmd.Stage)
	switch stage {
	case rcmd.TrainStage:
		tableName = "ratings_train"
	case rcmd.PredictStage:
		tableName = "ratings_test"
	default:
		panic("unknown stage")
	}

	if stage == rcmd.TrainStage && recSys.ubcTrain != nil {
		cachedSeq, er := recSys.ubcTrain.Get(int64(userId), maxTs, maxLen)
		if er != nil {
			log.Debugf("failed to get user %d behavior from cache: %v", userId, er)
		} else {
			itemSeq = utils.Int64SeqToIntSeq(cachedSeq.Items)
			return
		}
	} else if stage == rcmd.PredictStage && recSys.ubcPredict != nil {
		cachedSeq, er := recSys.ubcPredict.Get(int64(userId), maxTs, maxLen)
		if er != nil {
			log.Debugf("failed to get user %d behavior from cache: %v", userId, er)
		} else {
			itemSeq = utils.Int64SeqToIntSeq(cachedSeq.Items)
			return
		}
	}

	rows, err = db.Query(`select movieId from `+tableName+
		` where userId = ? and timestamp <= ? order by timestamp desc limit ?`,
		userId, maxTs, maxLen)
	if err != nil {
		log.Errorf("failed to query ratings: %v", err)
		return
	}
	defer rows.Close()
	for rows.Next() {
		var movieId int
		if err = rows.Scan(&movieId); err != nil {
			log.Errorf("failed to scan movieId: %v", err)
			return
		}
		itemSeq = append(itemSeq, movieId)
	}

	return
}

//PreFillUbCache prefill ubcache with data from db `ub_test` or `ub_train`.
// the ub_train table is generated by SQL like:
// 	```sql
//	create table ratings_train_desc as
//		select r.userId, movieId, rating, timestamp
//			from ratings_train r order by r.userId, timestamp desc;
//
//  create table ub_train as
//		select userId, group_concat(movieId) movieIds ,group_concat(timestamp) timestamps
//			from ratings_train_desc group by userId order by timestamp;
//	```
// Sample from ub_train will be like:
// 	31699, "246,247,252,260,265", "825638410,825638407,825638403,825638401,825638400"
func PreFillUbCache(ubc *ubcache.UserBehaviorCache, table string) (err error) {
	rows, err := db.Query(`select userId, movieIds, timestamps from ` + table)
	if err != nil {
		log.Errorf("failed to query ratings: %v", err)
		return
	}
	defer rows.Close()
	for rows.Next() {
		var (
			userId     int
			movieIds   string
			timestamps string
			itemSeq    []int64
			tsSeq      []int64
		)
		if err = rows.Scan(&userId, &movieIds, &timestamps); err != nil {
			log.Errorf("failed to scan %s: %v", table, err)
			return
		}
		tsSeq = utils.ParseInt64Seq(timestamps)
		itemSeq = utils.ParseInt64Seq(movieIds)
		if len(itemSeq) != len(tsSeq) {
			log.Errorf("itemSeq and tsSeq length %d:%d not match", len(itemSeq), len(tsSeq))
			return
		}
		ubc.Set(int64(userId), &ubcache.TimeSeq{
			Ts:    tsSeq,
			Items: itemSeq,
		})
	}
	return
}

//PreRank is called before rank, it can be used to prefill ub cache.
func (recSys *MovielensRec) PreRank(ctx context.Context) (err error) {
	if recSys.ubcPredict == nil {
		recSys.ubcPredict = ubcache.NewUserBehaviorCache()
		err = PreFillUbCache(recSys.ubcPredict, "ub_test")
		if err != nil {
			log.Errorf("failed to prefill ubcache: %v", err)
			return
		}
	}
	return
}
