package main

import (
	"fmt"
	"sync/atomic"

	"github.com/pingcap/log"
	"go.uber.org/zap"
)

const batchProgressLogEvery = 20

type batchProgressTracker struct {
	totalFiles int64

	processedFiles  atomic.Int64
	downloadedBytes atomic.Int64
	parsedRecords   atomic.Int64
	uploadedBytes   atomic.Int64
	uploadedRecords atomic.Int64
	uploadParts     atomic.Int64
}

func newBatchProgressTracker(totalFiles int) *batchProgressTracker {
	return &batchProgressTracker{totalFiles: int64(totalFiles)}
}

func (p *batchProgressTracker) markFileProcessed(downloadBytes int64, records int, fileName string) {
	if p == nil {
		return
	}
	p.processedFiles.Add(1)
	p.downloadedBytes.Add(downloadBytes)
	p.parsedRecords.Add(int64(records))

	processed := p.processedFiles.Load()
	if processed == p.totalFiles || processed%batchProgressLogEvery == 0 {
		p.logProgress("batch progress", zap.String("lastFile", fileName))
	}
}

func (p *batchProgressTracker) markUploaded(partBytes, partRecords int64, filename string) {
	if p == nil {
		return
	}
	p.uploadedBytes.Add(partBytes)
	p.uploadedRecords.Add(partRecords)
	part := p.uploadParts.Add(1)
	p.logProgress("uploaded rolling output file",
		zap.String("filename", filename),
		zap.Int64("partBytes", partBytes),
		zap.Int64("partRecords", partRecords),
		zap.Int64("uploadPart", part),
	)
}

func (p *batchProgressTracker) logProgress(event string, extra ...zap.Field) {
	if p == nil {
		return
	}
	processed := p.processedFiles.Load()
	total := p.totalFiles
	var pct float64
	if total > 0 {
		pct = float64(processed) * 100 / float64(total)
	}
	fields := []zap.Field{
		zap.String("fileProgress", fmt.Sprintf("%d/%d", processed, total)),
		zap.Float64("fileProgressPct", pct),
		zap.Int64("downloadedBytes", p.downloadedBytes.Load()),
		zap.Int64("uploadedBytes", p.uploadedBytes.Load()),
		zap.Int64("parsedRecords", p.parsedRecords.Load()),
		zap.Int64("uploadedRecords", p.uploadedRecords.Load()),
		zap.Int64("uploadParts", p.uploadParts.Load()),
	}
	fields = append(fields, extra...)
	log.Info(event, fields...)
}