package store

import (
	"cmp"
	"context"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	duckdb "github.com/marcboeker/go-duckdb"
	"github.com/pingcap/log"
	dto "github.com/prometheus/client_model/go"
	"github.com/tidbcloud/serverless-metering/pkg/config"
	"github.com/tidbcloud/serverless-metering/pkg/monitor"
	"github.com/tidbcloud/serverless-metering/pkg/util"
	"go.uber.org/zap"
)

const (
	metricPrefix         = "metric"
	fileSuffix           = ".parquet"
	createMetricTableSQL = `CREATE TABLE IF NOT EXISTS %s (
	name varchar,
	type int64,
	labels JSON,
	timestamp int64,
	value int64
)`
)

type LocalParquetLogger struct {
	workDir        string
	maxBackups     int
	writeBatchSize int
	db             *sql.DB
	conn           driver.Conn
	metricCh       chan *Metric
	// When the table is exporting, we need to ensure partial file will not be upload to s3
	exportingTs atomic.Int64

	currentHour int64
	appender    *duckdb.Appender
}

func NewLocalParquetLogger() (LocalStorage, error) {
	cfg := config.Global().Store.Local
	if cfg.MaxRotateSeconds <= 0 {
		log.Info("disable local storage", zap.Int("maxRotateSeconds", cfg.MaxRotateSeconds))
		return &LocalParquetLogger{}, nil
	}

	workDir := filepath.Dir(cfg.Filename)
	err := util.EnsureDir(workDir)
	if err != nil {
		return nil, fmt.Errorf("failed to ensure dir: %s", err)
	}

	ctx := context.TODO()
	connector, err := duckdb.NewConnector(cfg.Filename, func(execerContext driver.ExecerContext) error {
		_, err := execerContext.ExecContext(ctx, fmt.Sprintf("SET memory_limit = '%dMiB'", cfg.MemoryLimit), nil)
		return err
	})
	if err != nil {
		return nil, fmt.Errorf("failed to create duckdb connector: %s", err)
	}

	conn, err := connector.Connect(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to duckdb: %s", err)
	}
	db := sql.OpenDB(connector)
	logger := LocalParquetLogger{
		writeBatchSize: cfg.WriteBatch,
		workDir:        workDir,
		maxBackups:     cfg.MaxBackups,
		db:             db,
		conn:           conn,
		metricCh:       make(chan *Metric, 100),
	}

	go logger.startBackgroundWriter()

	go func() {
		ticker := time.NewTicker(time.Second * time.Duration(cfg.TickSeconds))
		for {
			<-ticker.C
			logger.tick()
		}
	}()

	return &logger, nil
}

func (l *LocalParquetLogger) Log(labels map[string]string, metricFamilies map[string]*dto.MetricFamily) error {
	for _, mf := range metricFamilies {
		name := *mf.Name
		metricType := *mf.Type
		for _, metric := range mf.Metric {
			var value float64
			switch metricType {
			case dto.MetricType_GAUGE:
				value = *metric.Gauge.Value
			case dto.MetricType_COUNTER:
				value = *metric.Counter.Value
			default:
				// only support gauge and counter for now
				log.Warn("unsupported metric type", zap.Any("metric", metric))
				continue
			}
			if value <= 0 {
				// tikv storage has a lot of 0 values, skip them
				continue
			}

			lbl := make(map[string]string)
			for k, v := range labels {
				lbl[k] = v
			}
			for _, l := range metric.Label {
				lbl[*l.Name] = *l.Value
			}

			l.metricCh <- &Metric{
				Timestamp: 0, // Timestamp will determine by duckdb writer
				Name:      name,
				Type:      int32(metricType),
				Value:     value,
				Labels:    lbl,
			}
		}
	}

	return nil
}

func (l *LocalParquetLogger) LoadFiles(cursor int64) ([]*LocalFile, error) {
	return l.listFilesWithMinTs(cursor)
}

func (l *LocalParquetLogger) trySwitchTable(ctx context.Context, nowTs int64) error {
	currentHour := util.LocateLatestOnTheHour(nowTs)
	if currentHour <= l.currentHour {
		return nil
	}

	tableName := fmt.Sprintf("%s%d", metricPrefix, currentHour)
	_, err := l.db.ExecContext(ctx, fmt.Sprintf(createMetricTableSQL, tableName))
	if err != nil {
		return err
	}
	appender, err := duckdb.NewAppenderFromConn(l.conn, "", tableName)
	if err != nil {
		return err
	}

	if l.appender != nil {
		l.appender.Close()
	}
	l.appender = appender
	l.currentHour = currentHour
	monitor.LocalLoggerRotateProgressGauge.WithLabelValues("parquet").Set(float64(l.currentHour))
	return nil
}

// NOTE: write in single goroutine to avoid concurrecy issues.
func (l *LocalParquetLogger) startBackgroundWriter() {
	count := 0
	lastWriteTime := time.Now()
	for {
		ctx := context.TODO()
		metric := <-l.metricCh
		monitor.LocalLoggerWriteCounter.WithLabelValues("parquet", "received").Inc()
		nowTs := time.Now().Unix()
		_ = l.trySwitchTable(ctx, nowTs)

		labels, err := json.Marshal(metric.Labels)
		if err != nil {
			log.Warn("failed to marshal labels", zap.Any("lbl", metric.Labels), zap.Error(err))
			continue
		}

		err = l.appender.AppendRow(metric.Name, metric.Type, labels, nowTs, math.Round(metric.Value))
		if err != nil {
			log.Error("failed to append metrics", zap.Error(err))
			continue
		}

		count += 1
		// TODO: When the process restarts, data may loss if the appender hasn't flushed.
		// But currenly loss a few data doesn't influence the result. Optimize this in the future if needed.
		// TODO: aliyun data may small, maybe we need to flush more frequently.
		if count >= l.writeBatchSize || time.Since(lastWriteTime) > time.Minute {
			log.Debug("local parquet logger flush")
			err = l.appender.Flush()
			if err != nil {
				log.Error("failed to append metrics", zap.Error(err))
				continue
			}
			monitor.LocalLoggerWriteCounter.WithLabelValues("parquet", "flush").Add(float64(l.writeBatchSize))
			count = 0
			lastWriteTime = time.Now()
		}
	}
}

func (l *LocalParquetLogger) tick() {
	log.Debug("local parquet logger tick")
	// The tick metric will ensure metric table not empty even if the metric api is never called.
	l.metricCh <- &Metric{
		Timestamp: 0, // Timestamp will determine by duckdb writer
		Name:      "tick",
		Type:      0,
		Value:     0,
		Labels:    map[string]string{},
	}
	ctx := context.TODO()
	rows, err := l.db.QueryContext(ctx, "SHOW TABLES", nil)
	if err != nil {
		log.Error("failed to show tables", zap.Error(err))
		return
	}
	for rows.Next() {
		var tableName string
		err = rows.Scan(&tableName)
		if err != nil {
			log.Error("failed to read table name", zap.Error(err))
			continue
		}
		err = l.convertTableToParquet(tableName)
		if err != nil {
			log.Error("failed to convert table to parquet", zap.Error(err), zap.String("tableName", tableName))
			continue
		}
	}

	files, err := l.listFilesWithMinTs(0)
	if err != nil {
		log.Error("failed to list all files", zap.Error(err))
		return
	}
	if len(files) <= l.maxBackups {
		return
	}
	filesToDelete := files[:len(files)-l.maxBackups]
	for _, f := range filesToDelete {
		log.Info("deleting file", zap.String("fileName", f.Path))
		err = os.Remove(f.Path)
		if err != nil {
			log.Error("failed to delete file", zap.Error(err), zap.String("fileName", f.Path))
		}
	}
}

func (l *LocalParquetLogger) convertTableToParquet(tableName string) error {
	tableTs, err := l.parseTs(tableName)
	if err != nil {
		return err
	}
	if tableTs >= l.currentHour {
		return nil
	}

	log.Info("exporting table", zap.String("tableName", tableName))
	l.exportingTs.Store(tableTs)
	defer l.exportingTs.Store(0)
	dstFileName := fmt.Sprintf("%s/%s%s", l.workDir, tableName, fileSuffix)
	_, err = l.db.Exec(fmt.Sprintf("COPY %s TO '%s' (FORMAT PARQUET)", tableName, dstFileName))
	if err != nil {
		os.Remove(dstFileName)
		return fmt.Errorf("failed to export table to parquet: %s", err)
	}

	log.Info("deleting table", zap.String("tableName", tableName))
	_, err = l.db.Exec(fmt.Sprintf("DROP TABLE %s", tableName))
	if err != nil {
		return fmt.Errorf("failed to delete table: %s", err)
	}

	return nil
}

func (l *LocalParquetLogger) listFilesWithMinTs(minTs int64) ([]*LocalFile, error) {
	entries, err := os.ReadDir(l.workDir)
	if err != nil {
		return nil, fmt.Errorf("failed to read dir: %s", err)
	}
	files := make([]*LocalFile, 0)
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		entryName := entry.Name()
		if !strings.HasPrefix(entryName, metricPrefix) {
			continue
		}
		if !strings.HasSuffix(entryName, fileSuffix) {
			continue
		}
		fileTs, err := l.parseTs(entryName)
		if err != nil {
			return nil, err
		}
		if fileTs == l.exportingTs.Load() {
			continue
		}
		fileInfo, err := entry.Info()
		if err != nil {
			return nil, err
		}
		if fileInfo.Size() == 0 {
			continue
		}
		if fileTs <= minTs {
			continue
		}
		files = append(files, &LocalFile{
			Path: fmt.Sprintf("%s/%s", l.workDir, entryName),
			End:  fileTs + 3600 - 1,
		})
	}
	slices.SortFunc(files, func(i, j *LocalFile) int {
		return cmp.Compare(i.End, j.End)
	})
	return files, nil
}

func (l *LocalParquetLogger) parseTs(name string) (int64, error) {
	if !strings.HasPrefix(name, metricPrefix) {
		return 0, fmt.Errorf("invalid name: %s", name)
	}
	tsStr := strings.TrimSuffix(strings.TrimPrefix(name, metricPrefix), fileSuffix)
	ts, err := strconv.ParseInt(tsStr, 10, 64)
	if err != nil {
		return 0, fmt.Errorf("failed to parse ts: %s", err)
	}
	return ts, nil
}