package repository

import (
	"context"

	"github.com/tidbcloud/pkgv2/db"
	"gorm.io/gorm/clause"

	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/domain/repository"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/infra/repository/schema"
)

type zeroRentRecordRepositoryImpl struct {
	dbProvider db.Provider
}

func NewZeroRentRecordRepository(dbProvider db.Provider) repository.ZeroRentRecordRepository {
	return &zeroRentRecordRepositoryImpl{dbProvider: dbProvider}
}

func (r *zeroRentRecordRepositoryImpl) Upsert(
	ctx context.Context,
	clusterId uint64,
	encryptedSecret string,
	sourceIP string,
	tag string,
) error {
	ctx, span := tracer.Start(ctx, "ZeroRentRecordRepo.Upsert")
	defer span.End()

	wdb := r.dbProvider.WriteDB(ctx)

	record := &schema.ZeroRentRecord{
		ClusterID:       clusterId,
		EncryptedSecret: encryptedSecret,
		SourceIP:        sourceIP,
		Tag:             tag,
	}

	return wdb.Clauses(clause.OnConflict{
		Columns:   []clause.Column{{Name: "cluster_id"}},
		DoUpdates: clause.AssignmentColumns([]string{"encrypted_secret", "source_ip", "tag", "updated_at"}),
	}).Create(record).Error
}

func (r *zeroRentRecordRepositoryImpl) CountBySourceIP(
	ctx context.Context,
	limit int,
) ([]repository.SourceIPCount, error) {
	ctx, span := tracer.Start(ctx, "ZeroRentRecordRepo.CountBySourceIP")
	defer span.End()

	// Guard against non-positive limits to avoid unexpected or unbounded queries.
	if limit <= 0 {
		return []repository.SourceIPCount{}, nil
	}

	rdb := r.dbProvider.ReadDB(ctx)

	type row struct {
		SourceIP string `gorm:"column:source_ip"`
		Count    int64  `gorm:"column:count"`
	}

	var rows []row
	err := rdb.Model(&schema.ZeroRentRecord{}).
		Select("source_ip, COUNT(*) AS count").
		Group("source_ip").
		Order("count DESC").
		Limit(limit).
		Scan(&rows).Error
	if err != nil {
		return nil, err
	}

	result := make([]repository.SourceIPCount, 0, len(rows))
	for _, row := range rows {
		result = append(result, repository.SourceIPCount{
			SourceIP: row.SourceIP,
			Count:    row.Count,
		})
	}
	return result, nil
}

func (r *zeroRentRecordRepositoryImpl) GetByClusterID(
	ctx context.Context,
	clusterId uint64,
) (string, error) {
	ctx, span := tracer.Start(ctx, "ZeroRentRecordRepo.GetByClusterID")
	defer span.End()

	rdb := r.dbProvider.ReadDB(ctx)

	var record schema.ZeroRentRecord
	err := rdb.Model(&schema.ZeroRentRecord{}).
		Where("cluster_id = ?", clusterId).
		Take(&record).Error
	if err != nil {
		return "", err
	}
	return record.EncryptedSecret, nil
}