package store

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"math"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"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/util"
	"go.uber.org/zap"
)

const (
	createRcuMetricTableSQL = `CREATE TABLE IF NOT EXISTS rcu_metric (
  name varchar,
  type int64,
  labels JSON,
  timestamp int64,
  value int64
)`

	clearExpiredRcuMetricSQL = `DELETE FROM rcu_metric WHERE timestamp < ?`

	createRcuCapacityTableSQL = `CREATE TABLE IF NOT EXISTS rcu_capacity (
  keyspace_id int64 primary key,
  cluster_id string,
  min_rcu int64,
  max_rcu int64,
  provisioned_capacity int64,
  created_at timestamp default current_localtimestamp(),
  updated_at timestamp default current_localtimestamp()
)`

	writeRcuCapacitySQL = `
INSERT INTO
  rcu_capacity (
    keyspace_id,
	cluster_id,
    min_rcu,
    max_rcu,
    provisioned_capacity,
    updated_at
  )
VALUES
  (?, ?, ?, ?, ?, current_localtimestamp())
  ON CONFLICT DO UPDATE
SET
  cluster_id = EXCLUDED.cluster_id,
  min_rcu = EXCLUDED.min_rcu,
  max_rcu = EXCLUDED.max_rcu,
  provisioned_capacity = EXCLUDED.provisioned_capacity,
  updated_at = current_localtimestamp()`

	listRcuCapacitiesSQL = `SELECT
  keyspace_id,
  cluster_id,
  min_rcu,
  max_rcu,
  provisioned_capacity,
FROM
  rcu_capacity
`
	deleteRcuCapacitySQL = `DELETE FROM rcu_capacity WHERE keyspace_id = ?`

	calculateRcuUsageV1SQL = `SELECT
  keyspace_id,
  segment,
  sum(if (incr < 0, 0, incr)) AS value
FROM
  (
    SELECT
      labels ->> '$.name' AS keyspace_id,
      floor(timestamp / 60) AS segment,
      value - lag(value) OVER (
        PARTITION BY
          (name, labels)
        ORDER BY
          timestamp
      ) AS incr
    FROM
      rcu_metric
    WHERE
      labels ->> '$.name' IN (%s)
      AND timestamp >= ?
      AND timestamp < ?
      AND type = 0
      AND name IN (
        'resource_manager_resource_unit_read_request_unit_sum',
        'resource_manager_resource_unit_write_request_unit_sum',
        'resource_manager_resource_unit_sql_layer_request_unit',
        'resource_manager_resource_unit_sql_layer_request_unit_sum'
      )
  )
GROUP BY
  keyspace_id,
  segment
ORDER BY
  keyspace_id,
  segment
`

	calculateRcuUsageV2SQL = `SELECT
  keyspace_id,
  segment,
  sum(if (incr < 0, 0, incr)) AS value
FROM
  (
    SELECT
      labels ->> '$.name' AS keyspace_id,
      floor(timestamp / 60) AS segment,
      value - lag(value) OVER (
        PARTITION BY
          (name, labels)
        ORDER BY
          timestamp
      ) AS incr
    FROM
      rcu_metric
    WHERE
      labels ->> '$.name' IN (%s)
      AND timestamp >= ?
      AND timestamp < ?
      AND type = 0
      AND name = 'resource_manager_resource_unit_request_unit_v2_sum'
  )
GROUP BY
  keyspace_id,
  segment
ORDER BY
  keyspace_id,
  segment
`
)

type RcuStoreImpl struct {
	db             *sql.DB
	conn           driver.Conn
	appender       *duckdb.Appender
	metricCh       chan *Metric
	writeBatchSize int
}

func NewRcuStore() (RcuStore, error) {
	cfg := config.Global().Store.Rcu
	err := util.EnsureDir(cfg.Directory)
	if err != nil {
		return nil, err
	}
	ctx := context.TODO()
	connector, err := duckdb.NewConnector(filepath.Join(cfg.Directory, "rcu.db"), nil)
	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)
	_, err = db.ExecContext(ctx, createRcuMetricTableSQL)
	if err != nil {
		return nil, fmt.Errorf("failed to create rcu metric table: %s", err)
	}
	_, err = db.ExecContext(ctx, createRcuCapacityTableSQL)
	if err != nil {
		return nil, fmt.Errorf("failed to create capacity table: %s", err)
	}
	appender, err := duckdb.NewAppenderFromConn(conn, "", "rcu_metric")
	if err != nil {
		return nil, fmt.Errorf("failed to create appender: %s", err)
	}
	store := &RcuStoreImpl{
		db:             db,
		conn:           conn,
		writeBatchSize: cfg.WriteBatch,
		metricCh:       make(chan *Metric, 100),
		appender:       appender,
	}

	go store.startBackgroundWriter()

	return store, nil
}

func (s *RcuStoreImpl) startBackgroundWriter() {
	count := 0
	lastWriteTime := time.Now()
	ticker := time.NewTicker(time.Minute * 5)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			_, err := s.db.Exec(clearExpiredRcuMetricSQL, time.Now().Unix()-int64(60*60))
			if err != nil {
				log.Error("failed to clear expired rcu metric", zap.Error(err))
				continue
			}
			log.Info("cleared expired rcu metric")
		case metric := <-s.metricCh:
			labels, err := json.Marshal(metric.Labels)
			if err != nil {
				log.Warn("failed to marshal labels", zap.Any("lbl", metric.Labels), zap.Error(err))
				continue
			}

			ts := time.Now().Unix()
			err = s.appender.AppendRow(metric.Name, metric.Type, labels, ts, 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.
			if count >= s.writeBatchSize || time.Since(lastWriteTime) > time.Minute {
				log.Debug("rcu store flush")
				err = s.appender.Flush()
				if err != nil {
					log.Error("failed to append metrics", zap.Error(err))
					continue
				}
				count = 0
				lastWriteTime = time.Now()
			}
		}
	}
}

func (s *RcuStoreImpl) WriteMetrics(labels map[string]string, mf *dto.MetricFamily) error {
	name := *mf.Name
	metricType := *mf.Type
	for _, metric := range mf.Metric {
		var value float64
		switch metricType {
		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
		}

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

func (s *RcuStoreImpl) CalRcuUsages(keyspaceIds []uint32, startTs int64, endTs int64, useRUV2 bool) ([]*RcuUsage, error) {
	if len(keyspaceIds) == 0 {
		return []*RcuUsage{}, nil
	}

	keyspaceIdStrs := make([]string, 0)
	for _, id := range keyspaceIds {
		keyspaceIdStrs = append(keyspaceIdStrs, fmt.Sprintf("'%d'", id))
	}
	sqlTemplate := calculateRcuUsageV1SQL
	if useRUV2 {
		sqlTemplate = calculateRcuUsageV2SQL
	}
	query := fmt.Sprintf(sqlTemplate, strings.Join(keyspaceIdStrs, ","))
	rows, err := s.db.Query(query, startTs, endTs)
	log.Debug("calRcuUsages query", zap.String("sql", query), zap.Int64("startTs", startTs), zap.Int64("endTs", endTs), zap.Bool("useRUV2", useRUV2))
	if err != nil {
		return nil, fmt.Errorf("failed to calculate rcu usages: %s", err)
	}
	defer rows.Close()

	usages := make([]*RcuUsage, 0)
	var usage *RcuUsage
	var lastKeyspaceID string = "invalid"
	for rows.Next() {
		var keyspaceID string
		var segment float64
		var value sql.NullFloat64
		err = rows.Scan(&keyspaceID, &segment, &value)
		if err != nil {
			return nil, fmt.Errorf("failed to scan rcu usage: %s", err)
		}
		if !value.Valid {
			log.Debug("rcu usage is null", zap.String("keyspaceID", keyspaceID), zap.Float64("segment", segment))
			continue
		}
		// change segment to timestamp
		timestamp := segment * 60
		// ru per 60 seconds
		rcu := value.Float64 / 60
		if lastKeyspaceID != keyspaceID {
			lastKeyspaceID = keyspaceID
			usage = &RcuUsage{
				KeyspaceID: keyspaceID,
				Timestamps: []int64{int64(timestamp)},
				Values:     []int64{int64(rcu)},
			}
			usages = append(usages, usage)
		} else {
			usage.Timestamps = append(usage.Timestamps, int64(timestamp))
			usage.Values = append(usage.Values, int64(rcu))
		}
	}
	log.Debug("calRcuUsages result", zap.Int("usageCnt", len(usages)))

	return usages, nil
}

func (s *RcuStoreImpl) WriteRcuCapacity(capacity *RcuCapacity) error {
	_, err := s.db.Exec(writeRcuCapacitySQL, capacity.KeyspaceID, strconv.FormatUint(capacity.ClusterID, 10), capacity.MinRcu, capacity.MaxRcu, capacity.ProvisionedCapacity)
	if err != nil {
		return fmt.Errorf("failed to write rcu capacity: %s", err)
	}
	return nil
}

func (s *RcuStoreImpl) ListRcuCapacities() ([]*RcuCapacity, error) {
	rows, err := s.db.Query(listRcuCapacitiesSQL)
	if err != nil {
		return nil, fmt.Errorf("failed to list rcu capacities: %s", err)
	}
	defer rows.Close()

	capacities := make([]*RcuCapacity, 0)
	for rows.Next() {
		var capacity RcuCapacity
		err = rows.Scan(&capacity.KeyspaceID, &capacity.ClusterID, &capacity.MinRcu, &capacity.MaxRcu, &capacity.ProvisionedCapacity)
		if err != nil {
			return nil, fmt.Errorf("failed to scan rcu capacity: %s", err)
		}
		capacities = append(capacities, &capacity)
	}

	return capacities, nil
}

func (s *RcuStoreImpl) DeleteRcuCapacity(keyspaceID uint32) error {
	_, err := s.db.Exec(deleteRcuCapacitySQL, keyspaceID)
	if err != nil {
		return fmt.Errorf("failed to delete rcu capacity: %s", err)
	}
	return nil
}