package main

import (
	"compress/gzip"
	"context"
	"encoding/csv"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

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

	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/storage"
)

const (
	defaultMaxOutputFileSize = 500 * 1024 * 1024 // 500MB compressed output
	defaultParallelWorkers   = 5
	rotateCheckInterval      = 100
)

func slowLogCSVHeader() []string {
	return []string{
		"Time", "Instance", "User", "Host", "ConnID", "SessionAlias",
		"ExecRetryCount", "ExecRetryTime", "QueryTime", "ParseTime",
		"CompileTime", "RewriteTime", "PreprocSubqueries", "PreprocSubqueriesTime",
		"OptimizeTime", "WaitTS", "PrewriteTime", "WaitPrewriteBinlogTime",
		"CommitTime", "GetCommitTsTime", "CommitBackoffTime", "BackoffTypes",
		"ResolveLockTime", "LocalLatchWaitTime", "WriteKeys", "WriteSize",
		"PrewriteRegion", "TxnRetry", "CopTime", "ProcessTime", "WaitTime",
		"BackoffTime", "LockKeysTime", "RequestCount", "TotalKeys", "ProcessKeys",
		"RocksdbDeleteSkippedCount", "RocksdbKeySkippedCount", "RocksdbBlockCacheHitCount",
		"RocksdbBlockReadCount", "RocksdbBlockReadByte", "DB", "IndexNames",
		"IsInternal", "Digest", "Stats", "CopProcAvg", "CopProcP90", "CopProcMax",
		"CopProcAddr", "CopWaitAvg", "CopWaitP90", "CopWaitMax", "CopWaitAddr",
		"MemMax", "DiskMax", "KVTotal", "PDTotal", "BackoffTotal",
		"WriteSQLResponseTotal", "ResultRows", "Warnings", "BackoffDetail",
		"Prepared", "Succ", "IsExplicitTxn", "IsWriteCacheTable", "PlanFromCache",
		"PlanFromBinding", "HasMoreResults", "ResourceGroup", "RequestUnitRead",
		"RequestUnitWrite", "TimeQueuedByRC", "Plan", "PlanDigest", "BinaryPlan",
		"PrevStmt", "Query", "ServerlessTenantID", "ServerlessProjectID",
		"ServerlessClusterID", "WorkerRole", "IsTiFlash", "TxnStartTs",
	}
}

func appendSlowLogCSVRow(writer *csv.Writer, log SlowLog) error {
	record := []string{
		log.Time.Format(time.RFC3339),
		log.Instance,
		log.User,
		log.Host,
		strconv.Itoa(log.ConnID),
		log.SessionAlias,
		strconv.Itoa(log.ExecRetryCount),
		strconv.FormatFloat(log.ExecRetryTime, 'f', -1, 64),
		strconv.FormatFloat(log.QueryTime, 'f', -1, 64),
		strconv.FormatFloat(log.ParseTime, 'f', -1, 64),
		strconv.FormatFloat(log.CompileTime, 'f', -1, 64),
		strconv.FormatFloat(log.RewriteTime, 'f', -1, 64),
		strconv.Itoa(log.PreprocSubqueries),
		strconv.FormatFloat(log.PreprocSubqueriesTime, 'f', -1, 64),
		strconv.FormatFloat(log.OptimizeTime, 'f', -1, 64),
		strconv.FormatFloat(log.WaitTS, 'f', -1, 64),
		strconv.FormatFloat(log.PrewriteTime, 'f', -1, 64),
		strconv.FormatFloat(log.WaitPrewriteBinlogTime, 'f', -1, 64),
		strconv.FormatFloat(log.CommitTime, 'f', -1, 64),
		strconv.FormatFloat(log.GetCommitTsTime, 'f', -1, 64),
		strconv.FormatFloat(log.CommitBackoffTime, 'f', -1, 64),
		log.BackoffTypes,
		strconv.FormatFloat(log.ResolveLockTime, 'f', -1, 64),
		strconv.FormatFloat(log.LocalLatchWaitTime, 'f', -1, 64),
		strconv.Itoa(log.WriteKeys),
		strconv.Itoa(log.WriteSize),
		strconv.Itoa(log.PrewriteRegion),
		strconv.Itoa(log.TxnRetry),
		strconv.FormatFloat(log.CopTime, 'f', -1, 64),
		strconv.FormatFloat(log.ProcessTime, 'f', -1, 64),
		strconv.FormatFloat(log.WaitTime, 'f', -1, 64),
		strconv.FormatFloat(log.BackoffTime, 'f', -1, 64),
		strconv.FormatFloat(log.LockKeysTime, 'f', -1, 64),
		strconv.Itoa(log.RequestCount),
		strconv.Itoa(log.TotalKeys),
		strconv.Itoa(log.ProcessKeys),
		strconv.Itoa(log.RocksdbDeleteSkippedCount),
		strconv.Itoa(log.RocksdbKeySkippedCount),
		strconv.Itoa(log.RocksdbBlockCacheHitCount),
		strconv.Itoa(log.RocksdbBlockReadCount),
		strconv.Itoa(log.RocksdbBlockReadByte),
		log.DB,
		log.IndexNames,
		strconv.FormatBool(log.IsInternal),
		log.Digest,
		log.Stats,
		strconv.FormatFloat(log.CopProcAvg, 'f', -1, 64),
		strconv.FormatFloat(log.CopProcP90, 'f', -1, 64),
		strconv.FormatFloat(log.CopProcMax, 'f', -1, 64),
		log.CopProcAddr,
		strconv.FormatFloat(log.CopWaitAvg, 'f', -1, 64),
		strconv.FormatFloat(log.CopWaitP90, 'f', -1, 64),
		strconv.FormatFloat(log.CopWaitMax, 'f', -1, 64),
		log.CopWaitAddr,
		strconv.Itoa(log.MemMax),
		strconv.Itoa(log.DiskMax),
		strconv.FormatFloat(log.KVTotal, 'f', -1, 64),
		strconv.FormatFloat(log.PDTotal, 'f', -1, 64),
		strconv.FormatFloat(log.BackoffTotal, 'f', -1, 64),
		strconv.FormatFloat(log.WriteSQLResponseTotal, 'f', -1, 64),
		strconv.Itoa(log.ResultRows),
		log.Warnings,
		log.BackoffDetail,
		strconv.FormatBool(log.Prepared),
		strconv.FormatBool(log.Succ),
		strconv.FormatBool(log.IsExplicitTxn),
		strconv.FormatBool(log.IsWriteCacheTable),
		strconv.FormatBool(log.PlanFromCache),
		strconv.FormatBool(log.PlanFromBinding),
		strconv.FormatBool(log.HasMoreResults),
		log.ResourceGroup,
		strconv.FormatFloat(log.RequestUnitRead, 'f', -1, 64),
		strconv.FormatFloat(log.RequestUnitWrite, 'f', -1, 64),
		strconv.FormatFloat(log.TimeQueuedByRC, 'f', -1, 64),
		log.Plan,
		log.PlanDigest,
		log.BinaryPlan,
		log.PrevStmt,
		log.Query,
		log.ServerlessTenantID,
		log.ServerlessProjectID,
		log.ServerlessClusterID,
		log.WorkerRole,
		strconv.FormatBool(log.IsTiFlash),
		strconv.FormatUint(log.TxnStartTs, 10),
	}
	return writer.Write(record)
}

type rollingGzipUploader struct {
	mu             sync.Mutex
	ctx            context.Context
	targetStorage  storage.ListStorage
	targetBucket   string
	targetFilePath string
	maxSize        int64
	namePrefix     string
	progress       *batchProgressTracker

	gzFile    *os.File
	gzPath    string
	gzWriter  *gzip.Writer
	csvWriter *csv.Writer

	recordCount int64
	startTime   time.Time
	endTime     time.Time
	hasRecords  bool
}

func newRollingGzipUploader(
	ctx context.Context,
	targetStorage storage.ListStorage,
	targetBucket string,
	targetFilePath string,
	maxSize int64,
	namePrefix string,
) *rollingGzipUploader {
	if maxSize <= 0 {
		maxSize = defaultMaxOutputFileSize
	}
	if namePrefix == "" {
		namePrefix = logKindSlowlog
	}
	return &rollingGzipUploader{
		ctx:            ctx,
		targetStorage:  targetStorage,
		targetBucket:   targetBucket,
		targetFilePath: targetFilePath,
		maxSize:        maxSize,
		namePrefix:     namePrefix,
	}
}

func (r *rollingGzipUploader) setProgressTracker(progress *batchProgressTracker) {
	r.progress = progress
}

func (r *rollingGzipUploader) AppendStatement(raw []byte, eventTime time.Time) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.gzFile == nil {
		if err := r.openNewFileLocked(); err != nil {
			return err
		}
	}

	row, err := statementJSONToCSVRow(raw, enableRuV2)
	if err != nil {
		return err
	}
	if err := r.csvWriter.Write(row); err != nil {
		return err
	}
	r.recordCount++
	if !r.hasRecords {
		r.hasRecords = true
	}
	if isValidEventTime(eventTime) {
		if r.startTime.IsZero() || eventTime.Before(r.startTime) {
			r.startTime = eventTime
		}
		if r.endTime.IsZero() || eventTime.After(r.endTime) {
			r.endTime = eventTime
		}
	}

	if r.recordCount%rotateCheckInterval == 0 {
		return r.maybeRotateLocked()
	}
	return nil
}

func (r *rollingGzipUploader) Append(log SlowLog) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.gzFile == nil {
		if err := r.openNewFileLocked(); err != nil {
			return err
		}
	}

	if err := appendSlowLogCSVRow(r.csvWriter, log); err != nil {
		return err
	}
	r.recordCount++
	if !r.hasRecords {
		r.hasRecords = true
	}
	if isValidEventTime(log.Time) {
		if r.startTime.IsZero() || log.Time.Before(r.startTime) {
			r.startTime = log.Time
		}
		if r.endTime.IsZero() || log.Time.After(r.endTime) {
			r.endTime = log.Time
		}
	}

	if r.recordCount%rotateCheckInterval == 0 {
		return r.maybeRotateLocked()
	}
	return nil
}

func (r *rollingGzipUploader) FlushPending() error {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.maybeRotateLocked()
}

func (r *rollingGzipUploader) Finish() error {
	r.mu.Lock()
	defer r.mu.Unlock()
	if !r.hasRecords {
		return nil
	}
	return r.uploadCurrentLocked()
}

func (r *rollingGzipUploader) maybeRotateLocked() error {
	if !r.hasRecords {
		return nil
	}
	if err := r.flushWritersLocked(); err != nil {
		return err
	}
	stat, err := r.gzFile.Stat()
	if err != nil {
		return err
	}
	if stat.Size() < r.maxSize {
		return nil
	}
	return r.uploadCurrentLocked()
}

func (r *rollingGzipUploader) openNewFileLocked() error {
	tmpDir := os.TempDir()
	r.gzPath = filepath.Join(tmpDir, fmt.Sprintf("%s_roll_%d.gz", r.namePrefix, time.Now().UnixNano()))
	gzFile, err := os.Create(r.gzPath)
	if err != nil {
		return fmt.Errorf("create rolling gzip file: %w", err)
	}
	r.gzFile = gzFile
	r.gzWriter = gzip.NewWriter(gzFile)
	r.csvWriter = csv.NewWriter(r.gzWriter)
	header := slowLogCSVHeader()
	if r.namePrefix == logKindStatement {
		header = statementCSVHeader(enableRuV2)
	}
	if err := r.csvWriter.Write(header); err != nil {
		return fmt.Errorf("write csv header: %w", err)
	}
	r.recordCount = 0
	r.hasRecords = false
	return nil
}

func (r *rollingGzipUploader) flushWritersLocked() error {
	if r.csvWriter != nil {
		r.csvWriter.Flush()
		if err := r.csvWriter.Error(); err != nil {
			return err
		}
	}
	if r.gzWriter != nil {
		return r.gzWriter.Flush()
	}
	return nil
}

func (r *rollingGzipUploader) uploadCurrentLocked() error {
	if !r.hasRecords || r.gzFile == nil {
		return nil
	}
	if err := r.flushWritersLocked(); err != nil {
		return err
	}
	if err := r.gzWriter.Close(); err != nil {
		return fmt.Errorf("close gzip writer: %w", err)
	}
	if err := r.gzFile.Close(); err != nil {
		return fmt.Errorf("close gzip file: %w", err)
	}

	startTimeStr := r.startTime.UTC().Format("20060102150405")
	endTimeStr := r.endTime.UTC().Format("20060102150405")
	md5HashGz, err := calculateFileMD5(r.gzPath)
	if err != nil {
		return fmt.Errorf("calculate md5: %w", err)
	}
	filename := fmt.Sprintf("%s_%s_%s_%s.gz", r.namePrefix, startTimeStr, endTimeStr, md5HashGz[:8])

	f, err := os.Open(r.gzPath)
	if err != nil {
		return fmt.Errorf("open gzip file for upload: %w", err)
	}
	defer f.Close()

	stat, err := f.Stat()
	if err != nil {
		return fmt.Errorf("stat gzip file: %w", err)
	}

	fullPath := filepath.Join(r.targetFilePath, filename)
	if err := r.targetStorage.SaveObject(r.ctx, r.targetBucket, fullPath, f); err != nil {
		return fmt.Errorf("upload rolling file: %w", err)
	}

	partRecords := r.recordCount
	partBytes := stat.Size()
	if r.progress != nil {
		r.progress.markUploaded(partBytes, partRecords, filename)
	} else {
		log.Info("uploaded rolling output file",
			zap.String("filename", filename),
			zap.Int64("compressedBytes", partBytes),
			zap.Int64("records", partRecords),
			zap.String("startTime", startTimeStr),
			zap.String("endTime", endTimeStr))
	}
	recordsWritten.Add(float64(partRecords))

	if err := os.Remove(r.gzPath); err != nil && !os.IsNotExist(err) {
		log.Warn("failed to remove rolling gzip temp file", zap.String("path", r.gzPath), zap.Error(err))
	}

	r.gzFile = nil
	r.gzWriter = nil
	r.csvWriter = nil
	r.gzPath = ""
	r.recordCount = 0
	r.hasRecords = false
	return nil
}