// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.

package stats

import (
	"fmt"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	teaconst "github.com/TeaOSLab/EdgeDNS/internal/const"
	"github.com/TeaOSLab/EdgeDNS/internal/events"
	"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
	"github.com/TeaOSLab/EdgeDNS/internal/rpc"
	"github.com/iwind/TeaGo/Tea"
	timeutil "github.com/iwind/TeaGo/utils/time"
	"strconv"
	"sync"
	"time"
)

var SharedManager = NewStatManager()

func init() {
	if !teaconst.IsMain {
		return
	}

	events.On(events.EventStart, func() {
		SharedManager.Start()
	})
}

type StatManager struct {
	statMap map[string]*pb.NSRecordHourlyStat // record@hour => *NSRecordHourlyStat
	locker  sync.Mutex
	ticker  *time.Ticker
}

func NewStatManager() *StatManager {
	return &StatManager{
		statMap: map[string]*pb.NSRecordHourlyStat{},
	}
}

func (this *StatManager) Start() {
	this.ticker = time.NewTicker(5 * time.Minute)
	if Tea.IsTesting() {
		this.ticker = time.NewTicker(1 * time.Minute)
	}
	go func() {
		defer func() {
			if r := recover(); r != nil {
				remotelogs.Error("STAT", fmt.Sprintf("goroutine panic: %v", r))
			}
		}()

		for range this.ticker.C {
			err := this.Loop()
			if err != nil {
				if rpc.IsConnError(err) {
					remotelogs.Debug("STAT", "upload failed: "+err.Error())
				} else {
					remotelogs.Error("STAT", "upload failed: "+err.Error())
				}
			}
		}
	}()
}

func (this *StatManager) Add(domainId int64, recordId int64, bytes int64) {
	this.locker.Lock()
	defer this.locker.Unlock()

	var hour = timeutil.Format("H")
	var key = strconv.FormatInt(recordId, 10) + "@" + hour
	stat, ok := this.statMap[key]
	if ok {
		stat.CountRequests++
		stat.Bytes += bytes
	} else {
		this.statMap[key] = &pb.NSRecordHourlyStat{
			NsDomainId:    domainId,
			NsRecordId:    recordId,
			Bytes:         bytes,
			CountRequests: 1,
			CreatedAt:     time.Now().Unix(),
		}
	}
}

func (this *StatManager) Loop() error {
	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return err
	}

	this.locker.Lock()
	var m = this.statMap
	this.statMap = map[string]*pb.NSRecordHourlyStat{}
	this.locker.Unlock()

	if len(m) == 0 {
		return nil
	}

	var stats = []*pb.NSRecordHourlyStat{}
	for _, stat := range m {
		stats = append(stats, stat)
	}

	_, err = rpcClient.NSRecordHourlyStatRPC.UploadNSRecordHourlyStats(rpcClient.Context(), &pb.UploadNSRecordHourlyStatsRequest{
		Stats: stats,
	})
	if err != nil {
		return err
	}
	return nil
}