package lease

import (
	"context"
	"errors"
	"fmt"
	"time"

	gomysql "github.com/go-sql-driver/mysql"
	"github.com/tidbcloud/pkgv2/db"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

const TableName = "distributed_leases"

const acquireMaxRetry = 5

var (
	ErrInvalidDB    = errors.New("lease locker db is nil")
	ErrInvalidKey   = errors.New("lease key is required")
	ErrInvalidOwner = errors.New("lease owner is required")
	ErrInvalidTTL   = errors.New("lease ttl must be greater than zero")
)

type Lease struct {
	LockKey    string    `gorm:"type:VARCHAR(512);primaryKey;not null"`
	OwnerID    string    `gorm:"type:VARCHAR(255);not null;index:idx_owner_lease_until,priority:1"`
	LeaseUntil time.Time `gorm:"type:DATETIME(3);not null;index:idx_owner_lease_until,priority:2"`
	CreatedAt  time.Time `gorm:"type:DATETIME(3);not null"`
	UpdatedAt  time.Time `gorm:"type:DATETIME(3);not null"`
}

func (Lease) TableName() string {
	return TableName
}

type Locker interface {
	Acquire(ctx context.Context, key, owner string, ttl time.Duration) (bool, error)
	Renew(ctx context.Context, key, owner string, ttl time.Duration) (bool, error)
	Release(ctx context.Context, key, owner string) (bool, error)
}

type DBLocker struct {
	provider db.Provider
}

func NewDBLocker(provider db.Provider) (*DBLocker, error) {
	if provider == nil {
		return nil, ErrInvalidDB
	}

	return &DBLocker{provider: provider}, nil
}

func NewLocker(provider db.Provider) (Locker, error) {
	return NewDBLocker(provider)
}

func (l *DBLocker) Acquire(ctx context.Context, key, owner string, ttl time.Duration) (bool, error) {
	if err := validateLeaseInput(key, owner, ttl); err != nil {
		return false, err
	}

	for attempt := range acquireMaxRetry {
		now := time.Now().UTC()
		leaseUntil := now.Add(ttl)
		acquired := false

		err := l.provider.WriteDB(ctx).WithContext(ctx).Transaction(func(tx *gorm.DB) error {
			var existing Lease
			err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
				Where("lock_key = ?", key).
				First(&existing).Error
			if errors.Is(err, gorm.ErrRecordNotFound) {
				newLease := &Lease{
					LockKey:    key,
					OwnerID:    owner,
					LeaseUntil: leaseUntil,
				}
				if createErr := tx.Create(newLease).Error; createErr != nil {
					if isDuplicateKeyError(createErr) {
						acquired = false
						return nil
					}
					return createErr
				}
				acquired = true
				return nil
			}
			if err != nil {
				return err
			}

			if existing.OwnerID != owner && existing.LeaseUntil.After(now) {
				acquired = false
				return nil
			}

			if updateErr := tx.Model(&Lease{}).
				Where("lock_key = ?", key).
				Updates(map[string]any{
					"owner_id":    owner,
					"lease_until": leaseUntil,
				}).Error; updateErr != nil {
				return updateErr
			}

			acquired = true
			return nil
		})
		if err == nil {
			return acquired, nil
		}
		if !isRetryableTxError(err) || attempt == acquireMaxRetry-1 {
			return false, fmt.Errorf("acquire lease failed: %w", err)
		}

		if sleepErr := sleepWithContext(ctx, time.Duration(attempt+1)*10*time.Millisecond); sleepErr != nil {
			return false, fmt.Errorf("acquire lease failed: %w", sleepErr)
		}
	}

	return false, nil
}

func (l *DBLocker) Renew(ctx context.Context, key, owner string, ttl time.Duration) (bool, error) {
	if err := validateLeaseInput(key, owner, ttl); err != nil {
		return false, err
	}

	now := time.Now().UTC()
	result := l.provider.WriteDB(ctx).WithContext(ctx).
		Model(&Lease{}).
		Where("lock_key = ? AND owner_id = ? AND lease_until > ?", key, owner, now).
		Updates(map[string]any{
			"lease_until": now.Add(ttl),
		})
	if result.Error != nil {
		return false, fmt.Errorf("renew lease failed: %w", result.Error)
	}

	return result.RowsAffected == 1, nil
}

func (l *DBLocker) Release(ctx context.Context, key, owner string) (bool, error) {
	if key == "" {
		return false, ErrInvalidKey
	}
	if owner == "" {
		return false, ErrInvalidOwner
	}

	result := l.provider.WriteDB(ctx).WithContext(ctx).
		Where("lock_key = ? AND owner_id = ?", key, owner).
		Delete(&Lease{})
	if result.Error != nil {
		return false, fmt.Errorf("release lease failed: %w", result.Error)
	}

	return result.RowsAffected == 1, nil
}

func validateLeaseInput(key, owner string, ttl time.Duration) error {
	if key == "" {
		return ErrInvalidKey
	}
	if owner == "" {
		return ErrInvalidOwner
	}
	if ttl <= 0 {
		return ErrInvalidTTL
	}
	return nil
}

func isDuplicateKeyError(err error) bool {
	var mysqlErr *gomysql.MySQLError
	if errors.As(err, &mysqlErr) {
		return mysqlErr.Number == 1062
	}
	return false
}

func isRetryableTxError(err error) bool {
	var mysqlErr *gomysql.MySQLError
	if errors.As(err, &mysqlErr) {
		return mysqlErr.Number == 1213 || mysqlErr.Number == 1205
	}
	return false
}

func sleepWithContext(ctx context.Context, d time.Duration) error {
	timer := time.NewTimer(d)
	defer timer.Stop()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}