go-ohif-proxy/internal/api/repository/shortcode.go

package repository

import (
	"fmt"
	"math/rand"
	"time"

	"devone.aplikasi.web.id/gitea/mario/go-ohif-proxy/internal/database"
	"github.com/jmoiron/sqlx"
)

// DBShortCode represents a shortcode from the database
type DBShortCode struct {
	ID                            int       `db:"id"`
	Shortcode                     string    `db:"Shortcode"`
	ShortcodeIsUsed               string    `db:"ShortcodeIsUsed"`
	ShortcodeIsUsedBy_ShortlinkID *int      `db:"ShortcodeIsUsedBy_ShortlinkID"`
	ShortcodeCreatedAt            time.Time `db:"ShortcodeCreatedAt"`
	ShortcodeUpdatedAt            time.Time `db:"ShortcodeUpdatedAt"`
}

// ShortCodeRepository handles database operations related to shortcodes
type ShortCodeRepository struct {
	*Repository
}

// NewShortCodeRepository creates a new shortcode repository
func NewShortCodeRepository() *ShortCodeRepository {
	return &ShortCodeRepository{
		Repository: NewRepository(),
	}
}

// GetUnusedShortCode retrieves an unused shortcode from the database
func (r *ShortCodeRepository) GetUnusedShortCode() (*string, error) {
	var shortcode string

	query := `SELECT Shortcode FROM shortcodes 
		WHERE ShortcodeIsUsed = 'N' 
		ORDER BY RAND() 
		LIMIT 1`

	err := database.DB.Get(&shortcode, query)
	if err != nil {
		return nil, fmt.Errorf("database error getting unused shortcode: %w", err)
	}

	return &shortcode, nil
}

// MarkShortCodeAsUsed marks a shortcode as used by a specific shortlink
func (r *ShortCodeRepository) MarkShortCodeAsUsed(tx *sqlx.Tx, shortcode string, shortlinkID int) error {
	query := `UPDATE shortcodes 
		SET ShortcodeIsUsed = 'Y', 
		ShortcodeIsUsedBy_ShortlinkID = ? 
		WHERE Shortcode = ?`

	_, err := tx.Exec(query, shortlinkID, shortcode)
	if err != nil {
		return fmt.Errorf("database error marking shortcode as used: %w", err)
	}

	return nil
}

// MarkShortCodeAsUnused marks a shortcode as unused when shortlink expires or is deleted
func (r *ShortCodeRepository) MarkShortCodeAsUnused(shortcode string) error {
	query := `UPDATE shortcodes 
		SET ShortcodeIsUsed = 'N', 
		ShortcodeIsUsedBy_ShortlinkID = NULL 
		WHERE Shortcode = ?`

	_, err := database.DB.Exec(query, shortcode)
	if err != nil {
		return fmt.Errorf("database error marking shortcode as unused: %w", err)
	}

	return nil
}

// MarkShortCodeAsUnusedByShortlinkID marks a shortcode as unused by shortlink ID
func (r *ShortCodeRepository) MarkShortCodeAsUnusedByShortlinkID(shortlinkID int) error {
	query := `UPDATE shortcodes 
		SET ShortcodeIsUsed = 'N', 
		ShortcodeIsUsedBy_ShortlinkID = NULL 
		WHERE ShortcodeIsUsedBy_ShortlinkID = ?`

	_, err := database.DB.Exec(query, shortlinkID)
	if err != nil {
		return fmt.Errorf("database error marking shortcode as unused by shortlink ID: %w", err)
	}

	return nil
}

// MarkExpiredShortlinkShortcodesAsUnused finds all shortcodes associated with expired shortlinks
// and marks them as unused
func (r *ShortCodeRepository) MarkExpiredShortlinkShortcodesAsUnused() (int, error) {
	query := `
		UPDATE shortcodes sc
		INNER JOIN shortlink sl ON sc.ShortcodeIsUsedBy_ShortlinkID = sl.ShortlinkID
		SET sc.ShortcodeIsUsed = 'N',
			sc.ShortcodeIsUsedBy_ShortlinkID = NULL
		WHERE (sl.ShortlinkExpiredAt < NOW() OR sl.ShortlinkIsRevoked = TRUE)
			AND sc.ShortcodeIsUsed = 'Y'
	`

	result, err := database.DB.Exec(query)
	if err != nil {
		return 0, fmt.Errorf("database error freeing expired shortlink shortcodes: %w", err)
	}

	count, err := result.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("error getting affected rows count: %w", err)
	}

	return int(count), nil
}

// GenerateUniqueShortCode generates a unique 5-character capital letter shortcode
func GenerateUniqueShortCode() string {
	const charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	const codeLength = 5

	code := make([]byte, codeLength)
	for i := range code {
		code[i] = charset[rand.Intn(len(charset))]
	}

	return string(code)
}

// SeedShortCodes seeds the shortcodes table with n unique 5-character codes
func (r *ShortCodeRepository) SeedShortCodes(n int) error {
	// Initialize random seed
	rand.Seed(time.Now().UnixNano())

	// Use a transaction for better performance
	tx, err := database.DB.Beginx()
	if err != nil {
		return fmt.Errorf("failed to start transaction: %w", err)
	}

	// Set up deferred rollback that will be canceled if we commit
	defer func() {
		if tx != nil {
			tx.Rollback()
		}
	}()

	// Check how many shortcodes already exist
	var count int
	err = tx.Get(&count, "SELECT COUNT(*) FROM shortcodes")
	if err != nil {
		return fmt.Errorf("failed to count existing shortcodes: %w", err)
	}

	fmt.Printf("Found %d existing shortcodes, generating %d more\n", count, n)

	// Track already generated codes to avoid duplicates
	generatedCodes := make(map[string]bool)

	// Get existing codes to avoid duplicates
	existingCodes := []string{}
	err = tx.Select(&existingCodes, "SELECT Shortcode FROM shortcodes")
	if err != nil {
		return fmt.Errorf("failed to retrieve existing shortcodes: %w", err)
	}

	// Add existing codes to the map
	for _, code := range existingCodes {
		generatedCodes[code] = true
	}

	// Prepare the insert statement
	stmt, err := tx.Prepare("INSERT INTO shortcodes (Shortcode, ShortcodeIsUsed) VALUES (?, 'N')")
	if err != nil {
		return fmt.Errorf("failed to prepare statement: %w", err)
	}
	defer stmt.Close()

	// Generate and insert the shortcodes
	inserted := 0
	attempts := 0
	maxAttempts := n * 10 // Limiting attempts to avoid infinite loop

	for inserted < n && attempts < maxAttempts {
		attempts++
		code := GenerateUniqueShortCode()

		if !generatedCodes[code] {
			_, err := stmt.Exec(code)
			if err != nil {
				continue // Skip if insertion fails and try another code
			}

			generatedCodes[code] = true
			inserted++

			// Print progress every 100 codes
			if inserted%100 == 0 {
				fmt.Printf("Generated %d/%d codes\n", inserted, n)
			}
		}
	}

	// Commit the transaction
	if err = tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	// Clear the tx to prevent the deferred rollback
	tx = nil

	fmt.Printf("Successfully generated %d unique shortcodes\n", inserted)
	return nil
}