package query

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/pingcap/log"
	"github.com/tidbcloud/serverless-metering/pkg/config"
	"github.com/tidbcloud/serverless-metering/pkg/store"
	"github.com/tidbcloud/serverless-metering/pkg/util"
	"go.uber.org/zap"

	duckdb "github.com/marcboeker/go-duckdb"
)

type QueryOption struct {
	withClusterId bool
	withShard     bool
}

type DuckDBQueryEngine struct {
	db               *sql.DB
	workDir          string
	conn             driver.Conn
	metricTableNames []string
	usageTableName   string
	dbFilePath       string
	storage          store.Storage
	ts               int64
	offset           int
	enableRcu        bool
	shards           int // shard data by keyspace id, then query multiple times, used for OOM query.
}

func NewDuckDBQueryEngine(ctx context.Context, storage store.Storage, ts int64, offset int) (QueryEngine, error) {
	cfg := config.Global().Query
	err := util.EnsureDir(cfg.Directory)
	if err != nil {
		return nil, fmt.Errorf("failed to ensure dir: %s", err)
	}

	dbFilePath := filepath.Join(cfg.Directory, fmt.Sprintf("%d-%d.db", ts, offset))
	connector, err := duckdb.NewConnector(dbFilePath, 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)
	usageTableName := "usage"
	_, err = db.ExecContext(ctx, fmt.Sprintf(createUsageTableSQL, usageTableName, usageTableName))
	if err != nil {
		return nil, fmt.Errorf("failed to create usage table: %s", err)
	}

	engine := &DuckDBQueryEngine{
		db:             db,
		workDir:        cfg.Directory,
		conn:           conn,
		usageTableName: usageTableName,
		dbFilePath:     dbFilePath,
		storage:        storage,
		ts:             ts,
		offset:         offset,
		enableRcu:      config.Global().EnableRcu,
		shards:         cfg.Shards,
	}

	ft := store.GetFileType()
	switch ft {
	case store.FileTypeJson:
		err = engine.loadDataFromMetrics(ctx, ts, offset)
	case store.FileTypeParquet:
		err = engine.loadDataFromRemoteFiles(ctx, ts, offset)
	default:
		return nil, fmt.Errorf("unknown file type")
	}
	if err != nil {
		engine.Close()
		return nil, err
	}

	if err := engine.initLoadDataTaskTable(ctx); err != nil {
		engine.Close()
		return nil, err
	}

	return engine, nil
}

func (e *DuckDBQueryEngine) loadDataFromMetrics(ctx context.Context, ts int64, offset int) error {
	metricTableNames := make([]string, 0, offset)
	for i := 0; i < offset; i++ {
		scanTs := ts + int64(i*60*60)
		tableName := fmt.Sprintf("metric%d", scanTs)
		metricTableNames = append(metricTableNames, tableName)
		_, err := e.db.ExecContext(ctx, fmt.Sprintf(createMetricTableSQL, tableName))
		if err != nil {
			return fmt.Errorf("failed to create metric table: %s", err)
		}

		appender, err := duckdb.NewAppenderFromConn(e.conn, "", tableName)
		if err != nil {
			return fmt.Errorf("failed to create metric appender: %s", err)
		}
		defer appender.Close()

		scanner, err := e.storage.BuildRemoteScanner(ctx, scanTs)
		if err != nil {
			return fmt.Errorf("failed to build remote scanner: %s", err)
		}
		more := true
		for more {
			var metrics []*store.Metric
			metrics, more, err = scanner.ScanMetrics(ctx)
			if err != nil {
				return err
			}
			log.Debug("start write duckdb", zap.Int("size", len(metrics)))
			for _, metric := range metrics {
				labels, err := json.Marshal(metric.Labels)
				if err != nil {
					return err
				}
				err = appender.AppendRow(metric.Name, metric.Type, labels, metric.Timestamp, metric.Value)
				if err != nil {
					return err
				}
			}
			err = appender.Flush()
			if err != nil {
				return err
			}
			log.Debug("write duckdb finished", zap.Int("size", len(metrics)))
		}
	}

	e.metricTableNames = metricTableNames
	return nil
}

func (e *DuckDBQueryEngine) loadDataFromRemoteFiles(ctx context.Context, ts int64, offset int) error {
	metricTableNames := make([]string, 0, offset)
	for i := 0; i < offset; i++ {
		scanTs := ts + int64(i*60*60)
		scanner, err := e.storage.BuildRemoteScanner(ctx, scanTs)
		if err != nil {
			return fmt.Errorf("failed to build remote scanner: %s", err)
		}
		more := true
		tmpFileNames := make([]string, 0)
		for more {
			var remoteFile *store.RemoteFile
			remoteFile, more, err = scanner.ScanRemoteFiles(ctx)
			if err != nil {
				return err
			}
			if remoteFile == nil {
				continue
			}
			tmpFileName := filepath.Join(e.workDir, remoteFile.Name)

			err = func() error {
				defer remoteFile.Body.Close()
				data, err := io.ReadAll(remoteFile.Body)
				if err != nil {
					return fmt.Errorf("failed to read remote file: %s", err)
				}
				err = os.WriteFile(tmpFileName, data, 0777)
				if err != nil {
					return fmt.Errorf("failed to write file: %s", err)
				}
				tmpFileNames = append(tmpFileNames, tmpFileName)
				return nil
			}()
			if err != nil {
				return err
			}
		}
		if len(tmpFileNames) == 0 {
			continue
		}

		tableName := fmt.Sprintf("metric%d", scanTs)
		metricTableNames = append(metricTableNames, tableName)
		parquetNames := make([]string, len(tmpFileNames))
		for i, name := range tmpFileNames {
			parquetNames[i] = fmt.Sprintf("'%s'", name)
		}
		createTableSQL := fmt.Sprintf(
			"DROP TABLE IF EXISTS %s; CREATE TABLE %s AS SELECT * FROM read_parquet([%s])",
			tableName,
			tableName,
			strings.Join(parquetNames, ","),
		)
		log.Info("create table from parquet", zap.String("sql", createTableSQL))
		_, err = e.db.ExecContext(ctx, createTableSQL)
		if err != nil {
			return fmt.Errorf("failed to create metric table: %s", err)
		}

		for _, name := range tmpFileNames {
			log.Debug("deleting temp file", zap.String("filename", name))
			err = os.Remove(name)
			if err != nil {
				log.Error("failed to delete temp file", zap.String("filename", name), zap.Error(err))
			}
		}
	}

	e.metricTableNames = metricTableNames
	return nil
}

func (e *DuckDBQueryEngine) Close() error {
	logger := log.L().With(zap.String("dbFilePath", e.dbFilePath))
	err := e.conn.Close()
	if err != nil {
		logger.Error("close conn failed", zap.Error(err))
	}
	err = e.db.Close()
	if err != nil {
		logger.Error("close db failed", zap.Error(err))
	}
	err = os.Remove(fmt.Sprintf("%s.wal", e.dbFilePath))
	if err != nil {
		logger.Info("remove wal file failed", zap.Error(err))
	}
	err = os.Remove(fmt.Sprintf("%s.tmp", e.dbFilePath))
	if err != nil {
		logger.Info("remove tmp file failed", zap.Error(err))
	}
	err = os.Remove(e.dbFilePath)
	if err != nil {
		logger.Info("remove db file failed", zap.Error(err))
	}
	return nil
}

func (e *DuckDBQueryEngine) CalLoadDataRu(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, loadDataRuSQL, MetricTypeLoadDataRU, &QueryOption{})
}

func (e *DuckDBQueryEngine) CalPublicEgress(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, publicEgressSQL, MetricTypePublicEgress, &QueryOption{withClusterId: true})
}

func (e *DuckDBQueryEngine) CalPrivatelinkEgress(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, privatelinkEgressSQL, MetricTypePrivatelinkEgress, &QueryOption{withClusterId: true})
}

func (e *DuckDBQueryEngine) CalTidbRU(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tidbRuSQL, MetricTypeTidbRU, &QueryOption{withShard: true})
}

func (e *DuckDBQueryEngine) CalTiflashRU(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tiflashRuSQL, MetricTypeTiflashRU, &QueryOption{})
}

func (e *DuckDBQueryEngine) CalTidbRUV2(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tidbRUV2SQL, MetricTypeTidbRUV2, &QueryOption{withShard: true})
}

func (e *DuckDBQueryEngine) CalTiflashRUV2(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tiflashRUV2SQL, MetricTypeTiflashRUV2, &QueryOption{})
}

func (e *DuckDBQueryEngine) CalTikvStorageMedium(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tikvStorageMediumSQL, MetricTypeTikvStorage, &QueryOption{withShard: true})
}

func (e *DuckDBQueryEngine) CalTiflashStorageMedium(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tiflashStorageMediumSQL, MetricTypeTiflashStorage, &QueryOption{})
}

func (e *DuckDBQueryEngine) CalTikvStorageLarge(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tikvStorageLargeSQL, MetricTypeTikvStorage, &QueryOption{withShard: true})
}

func (e *DuckDBQueryEngine) CalTikvStorageIAMedium(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tikvStorageIAMediumSQL, MetricTypeTikvStorageIA, &QueryOption{withShard: true})
}

func (e *DuckDBQueryEngine) CalTikvStorageIALarge(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tikvStorageIALargeSQL, MetricTypeTikvStorageIA, &QueryOption{withShard: true})
}

func (e *DuckDBQueryEngine) CalTiflashStorageLarge(ctx context.Context) ([]*Usage, error) {
	return e.calculate(ctx, tiflashStorageLargeSQL, MetricTypeTiflashStorage, &QueryOption{})
}

func (e *DuckDBQueryEngine) CalIntermediateTikvStorageMedium(ctx context.Context) ([]*Usage, error) {
	return e.calculateIntermediate(ctx, tikvStorageMediumSQL, MetricTypeTikvStorage, &QueryOption{withShard: true})
}

func (e *DuckDBQueryEngine) CalIntermediateTikvStorageIAMedium(ctx context.Context) ([]*Usage, error) {
	return e.calculateIntermediate(ctx, tikvStorageIAMediumSQL, MetricTypeTikvStorageIA, &QueryOption{withShard: true})
}

func (e *DuckDBQueryEngine) CalIntermediateTiflashStorageMedium(ctx context.Context) ([]*Usage, error) {
	return e.calculateIntermediate(ctx, tiflashStorageMediumSQL, MetricTypeTiflashStorage, &QueryOption{})
}

func (e *DuckDBQueryEngine) CalRCU(ctx context.Context) ([]*Usage, error) {
	if !e.enableRcu {
		return []*Usage{}, nil
	}
	return e.calculate(ctx, rcuSQL, MetricTypeRCU, &QueryOption{})
}

func (e *DuckDBQueryEngine) CalExecCount(ctx context.Context) (ExecCount, error) {
	result := make(ExecCount)

	for _, tableName := range e.metricTableNames {
		rows, err := e.db.QueryContext(ctx, fmt.Sprintf(execCountSQL, tableName))
		if err != nil {
			log.Error("failed to query exec count", zap.Error(err))
			return nil, err
		}

		for rows.Next() {
			var db, clusterId string
			var value sql.NullInt64

			err = rows.Scan(&db, &clusterId, &value)
			if err != nil {
				log.Error("failed to scan exec count", zap.Error(err))
				return nil, err
			}

			if !value.Valid || value.Int64 <= 0 {
				continue
			}

			if result[clusterId] == nil {
				result[clusterId] = make(map[string]int64)
			}
			result[clusterId][db] += value.Int64
		}
		rows.Close()
	}

	return result, nil
}

func (e *DuckDBQueryEngine) calculateIntermediate(ctx context.Context, sql string, mt MetricType, option *QueryOption) ([]*Usage, error) {
	usages, err := e.calculate(ctx, sql, mt, option)
	if err != nil {
		return nil, err
	}
	data, _ := json.Marshal(usages)
	err = e.storage.PutIntermediateResult(ctx, e.ts, string(mt), data)
	if err != nil {
		return nil, err
	}
	log.Info("PutIntermediateResult success", zap.Int64("ts", e.ts), zap.String("digest", string(mt)))

	return []*Usage{}, nil
}

func (e *DuckDBQueryEngine) calculate(ctx context.Context, query string, metricType MetricType, option *QueryOption) ([]*Usage, error) {
	log.Info("start to calculate", zap.String("query", query))
	err := e.aggregateMetricsToUsage(ctx, query, metricType, option)
	if err != nil {
		return nil, fmt.Errorf("failed to aggregate metrics: %s", err)
	}
	usages, err := e.queryUsages(ctx, metricType, e.getQueryDigest(query))
	if err != nil {
		return nil, fmt.Errorf("failed to query usages: %s", err)
	}
	return usages, nil
}

func (e *DuckDBQueryEngine) queryUsages(ctx context.Context, metricType MetricType, digest string) ([]*Usage, error) {
	usages := make([]*Usage, 0)
	rows, err := e.db.QueryContext(ctx, fmt.Sprintf(queryUsageSQL, e.usageTableName, string(metricType), digest))
	if err != nil {
		return nil, err
	}
	for rows.Next() {
		var keyspaceId, primaryId string
		var value int64
		err = rows.Scan(&keyspaceId, &primaryId, &value)
		if err != nil {
			return nil, err
		}
		usages = append(usages, &Usage{
			KeyspaceId: keyspaceId,
			PrimaryId:  primaryId,
			Value:      value,
			MetricType: metricType,
		})
	}
	return usages, nil
}

func (e *DuckDBQueryEngine) aggregateMetricsToUsage(ctx context.Context, query string, metricType MetricType, option *QueryOption) error {
	digest := e.getQueryDigest(query)
	for _, tableName := range e.metricTableNames {
		// If withShard is true, we will query each shard separately.
		if option.withShard {
			for i := 0; i < e.shards; i++ {
				log.Debug("querying shard", zap.Int("shard", i))
				querySQL := fmt.Sprintf(query, tableName, e.shards, i)
				err := e.queryAndWriteUsageTable(ctx, querySQL, metricType, digest, option)
				if err != nil {
					log.Error("queryAndWriteUsageTable failed", zap.String("querySQL", querySQL), zap.Error(err), zap.Int("shard", i))
					return fmt.Errorf("failed to query and write usage table: %s", err)
				}
			}
		} else {
			querySQL := fmt.Sprintf(query, tableName)
			err := e.queryAndWriteUsageTable(ctx, querySQL, metricType, digest, option)
			if err != nil {
				log.Error("queryAndWriteUsageTable failed", zap.String("querySQL", querySQL), zap.Error(err))
				return fmt.Errorf("failed to query and write usage table: %s", err)
			}
		}
	}
	return nil
}

func (e *DuckDBQueryEngine) queryAndWriteUsageTable(ctx context.Context, querySQL string, metricType MetricType, digest string, option *QueryOption) error {
	rows, err := e.db.QueryContext(ctx, querySQL)
	if err != nil {
		return err
	}

	appender, err := duckdb.NewAppenderFromConn(e.conn, "", e.usageTableName)
	if err != nil {
		return err
	}
	defer appender.Close()
	for rows.Next() {
		var clusterId, keyspaceId string
		var result sql.NullFloat64

		if option.withClusterId {
			err = rows.Scan(&clusterId, &result)
		} else {
			err = rows.Scan(&keyspaceId, &result)
		}
		if err != nil {
			return err
		}
		if !result.Valid {
			log.Debug("skip null value", zap.String("clusterId", clusterId), zap.String("keyspaceId", keyspaceId))
			continue
		}
		value := result.Float64
		if value <= 0 {
			log.Debug("skip value less than zero", zap.String("clusterId", clusterId), zap.String("keyspaceId", keyspaceId))
			continue
		}

		err = appender.AppendRow(keyspaceId, clusterId, string(metricType), digest, int64(value))
		if err != nil {
			return err
		}
	}
	return nil
}

func (e *DuckDBQueryEngine) getQueryDigest(query string) string {
	return fmt.Sprintf("%x", sha256.Sum256([]byte(query)))
}

func (e *DuckDBQueryEngine) initLoadDataTaskTable(ctx context.Context) error {
	if _, err := e.db.ExecContext(ctx, createLoadDataTaskSQL); err != nil {
		return fmt.Errorf("failed create load data task table: %s", err)
	}

	filename := fmt.Sprintf("%s/load_data_task.parquet", e.workDir)
	if _, err := os.Stat(filename); err == nil {
		if _, err := e.db.ExecContext(ctx, fmt.Sprintf(importLoadDataTaskSQL, filename)); err != nil {
			return fmt.Errorf("failed to import load data task: %s", err)
		}
	}

	for _, metricTable := range e.metricTableNames {
		if _, err := e.db.ExecContext(ctx, fmt.Sprintf(writeLoadDataTaskSQL, metricTable)); err != nil {
			return fmt.Errorf("failed to write load data task: %s", err)
		}
	}

	if _, err := e.db.ExecContext(ctx, fmt.Sprintf(exportLoadDataTaskSQL, filename)); err != nil {
		return fmt.Errorf("failed to export load data task: %s", err)
	}

	return nil
}