package handler

import (
	"fmt"
	"log"
	"time"

	"github.com/go-pg/pg/v10"
	"github.com/go-pg/pg/v10/orm"
	jsoniter "github.com/json-iterator/go"
	"github.com/tal-tech/go-stash/stash/config"
	"github.com/tal-tech/go-stash/stash/filter"
	"github.com/tal-tech/go-zero/core/executors"
	"github.com/tal-tech/go-zero/core/threading"
)

const timestampKey = "time"

type MessageChunkHandlerPG struct {
	filters  []filter.FilterFunc
	db       *pg.DB
	conf     config.PostgresqlConf
	inserter *executors.ChunkExecutor
}

func NewMessageChunkHandlerPG(constant config.PostgresqlConf) *MessageChunkHandlerPG {
	DB := pg.Connect(&pg.Options{
		User:     constant.User,
		Password: constant.Password,
		Database: constant.DBName,
		Addr:     fmt.Sprintf("%s:%s", constant.Host, constant.Port),
	})
	handler := &MessageChunkHandlerPG{
		db:   DB,
		conf: constant,
	}
	handler.inserter = executors.NewChunkExecutor(handler.execute, executors.WithChunkBytes(constant.BatchSize*1024)) //大小 1KB * 100
	go handler.TimerCreateLogTable()
	return handler
}

func (mh *MessageChunkHandlerPG) AddFilters(filters ...filter.FilterFunc) {
	mh.filters = append(mh.filters, filters...)
}

func (mh *MessageChunkHandlerPG) Consume(_, val string) error {
	var m map[string]interface{}
	if err := jsoniter.Unmarshal([]byte(val), &m); err != nil {
		return err
	}

	for _, proc := range mh.filters {
		if m = proc(m); m == nil {
			return nil
		}
	}
	return mh.inserter.Add(m, len([]byte(val)))
}

func (mh *MessageChunkHandlerPG) execute(tasks []interface{}) {
	begin :=time.Now()
	var logs []*Logs
	for i := 0; i < len(tasks); i++ {
		item := tasks[i]
		if m, ok := item.(map[string]interface{}); ok {
			logs = append(logs, &Logs{
				Log:     m,
				LogTime: getTime(m),
			})
		}
	}
	if len(logs) > 0 {
		if _, err := mh.db.Model(&logs).Insert(); err != nil {
			fmt.Println("[logstash] Insert Error:", err)
		}

		end :=time.Now()
		fmt.Printf("[logstash] Insert task:%v cost:%v \n", len(tasks),end.Sub(begin))
	}
}

func (mh *MessageChunkHandlerPG) TimerCreateLogTable() {
	t := time.NewTimer(time.Hour * 6)
	fmt.Printf("[logstash] Begin TimerCreateLogTable \n")
	mh.createLogTable()
	for range t.C {
		threading.RunSafe(
			func() {
				mh.createLogTable()
			},
		)
	}
}

func (mh *MessageChunkHandlerPG) createLogTable() {
	var err error
	// creates database schema for Log models.
	err = mh.db.Model(&Logs{}).CreateTable(&orm.CreateTableOptions{
		IfNotExists: true,
	})
	if err != nil {
		log.Fatal(err)
	}

	logStartTime := time.Now()
	logEndTime := logStartTime.AddDate(0, 3, 0)
	for logStartTime.Unix() <= logEndTime.Unix() {
		// Before insert, always try create partition
		err = createNewPartition(mh.db, logStartTime)
		if err != nil {
			log.Fatal(err)
		}
		logStartTime = logStartTime.AddDate(0, 1, 0)
	}
}

func getTime(m map[string]interface{}) time.Time {
	if ti, ok := m[timestampKey]; ok {
		if ts, ok := ti.(string); ok {
			if t, err := time.Parse(time.RFC3339, ts); err == nil {
				return t
			}
		}
	}

	return time.Now()
}