package main

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"os"
	"runtime"
	"sync"
	"time"

	_ "github.com/lib/pq"
)

// ClipID represents a clip ID to update
type ClipID struct {
	ID string
}

// Worker pool configuration
type WorkerPool struct {
	numWorkers  int
	idsChan     chan []ClipID
	wg          sync.WaitGroup
	errorCount  int64
	updateCount int64
	mu          sync.Mutex
}

func main() {
	// Worker pool configuration
	numWorkers := getEnvInt("NUM_WORKERS", runtime.NumCPU())
	fetchBatchSize := getEnvInt("FETCH_BATCH_SIZE", 10000) // Fetch 10k IDs at a time
	updateBatchSize := getEnvInt("UPDATE_BATCH_SIZE", 500) // Update 500 records at a time per worker

	log.Printf("Starting Postgres cleanup with %d workers, fetch batch: %d, update batch: %d",
		numWorkers, fetchBatchSize, updateBatchSize)

	// Database connection
	dbURL := getEnv("DATABASE_URL", "postgres://suno:{DB_PASSWORD}@lyrics-backfill-cluster.cluster-cnfvffydbwvc.us-east-2.rds.amazonaws.com:5432/suno_main")

	// Connect to PostgreSQL (main connection for fetching IDs)
	mainDB, err := sql.Open("postgres", dbURL)
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}
	defer mainDB.Close()

	// Set connection pool parameters - need enough for all workers + main reader
	mainDB.SetMaxOpenConns(numWorkers + 2)
	mainDB.SetMaxIdleConns(numWorkers + 2)

	err = mainDB.Ping()
	if err != nil {
		log.Fatalf("Failed to ping database: %v", err)
	}
	log.Println("Successfully connected to PostgreSQL database")

	// Get approximate count for progress reporting
	var totalCount int64
	err = mainDB.QueryRow("SELECT reltuples::bigint AS estimate FROM pg_class WHERE relname = 'bots_generatedclip'").Scan(&totalCount)
	if err != nil {
		log.Printf("Warning: Could not get estimated count: %v", err)
		totalCount = 0
	} else {
		log.Printf("Estimated total records to process: %d", totalCount)
	}

	// Create worker pool
	pool := &WorkerPool{
		numWorkers:  numWorkers,
		idsChan:     make(chan []ClipID, numWorkers*2),
		errorCount:  0,
		updateCount: 0,
	}

	// Start worker goroutines - each worker gets its own DB connection
	for i := 0; i < pool.numWorkers; i++ {
		pool.wg.Add(1)
		go worker(i, pool, dbURL, updateBatchSize)
	}

	// Process in batches using ID-based cursor
	processedCount := 0
	lastID := ""
	startTime := time.Now()

	for {
		var query string
		var args []interface{}

		// Only select IDs where cleanup is needed (optimization)
		// Check if prompt_text is not empty OR if metadata->>'prompt' is not empty
		if lastID == "" {
			query = `SELECT id FROM bots_generatedclip 
					 WHERE (prompt_text IS NOT NULL AND prompt_text != '') 
					    OR (metadata->>'prompt' IS NOT NULL AND metadata->>'prompt' != '')
					 ORDER BY id LIMIT $1`
			args = []interface{}{fetchBatchSize}
		} else {
			query = `SELECT id FROM bots_generatedclip 
					 WHERE id > $1 
					   AND ((prompt_text IS NOT NULL AND prompt_text != '') 
					        OR (metadata->>'prompt' IS NOT NULL AND metadata->>'prompt' != ''))
					 ORDER BY id LIMIT $2`
			args = []interface{}{lastID, fetchBatchSize}
		}

		rows, err := mainDB.Query(query, args...)
		if err != nil {
			log.Fatalf("Failed to query database: %v", err)
		}

		batch := make([]ClipID, 0, updateBatchSize)
		recordCount := 0

		for rows.Next() {
			var clipID ClipID

			if err := rows.Scan(&clipID.ID); err != nil {
				log.Printf("Error scanning row: %v", err)
				continue
			}

			batch = append(batch, clipID)
			lastID = clipID.ID
			recordCount++

			// When we reach batch size, send to worker pool
			if len(batch) >= updateBatchSize {
				pool.idsChan <- batch
				batch = make([]ClipID, 0, updateBatchSize)
			}
		}

		rows.Close()

		// Send any remaining records in the batch
		if len(batch) > 0 {
			pool.idsChan <- batch
		}

		processedCount += recordCount

		// Report progress
		elapsed := time.Since(startTime)
		if totalCount > 0 {
			progress := float64(processedCount) / float64(totalCount) * 100
			recordsPerSec := float64(processedCount) / elapsed.Seconds()
			remainingRecords := totalCount - int64(processedCount)
			etaSeconds := float64(remainingRecords) / recordsPerSec

			log.Printf("Progress: %.2f%% (%d/%d) | Speed: %.0f rec/s | Updated: %d | Errors: %d | ETA: %s",
				progress, processedCount, totalCount, recordsPerSec,
				pool.updateCount, pool.errorCount,
				time.Duration(etaSeconds)*time.Second)
		} else {
			log.Printf("Processed %d records | Updated: %d | Errors: %d | Speed: %.0f rec/s",
				processedCount, pool.updateCount, pool.errorCount,
				float64(processedCount)/elapsed.Seconds())
		}

		// If we got fewer records than the batch size, we're done
		if recordCount < fetchBatchSize {
			break
		}
	}

	// Close the channel to signal workers to exit
	close(pool.idsChan)

	// Wait for all workers to finish
	pool.wg.Wait()

	elapsed := time.Since(startTime)
	log.Printf("Cleanup completed successfully in %s", elapsed)
	log.Printf("Total records scanned: %d", processedCount)
	log.Printf("Total records updated: %d", pool.updateCount)
	log.Printf("Total errors: %d", pool.errorCount)
}

// Worker function that processes batches of IDs
func worker(id int, pool *WorkerPool, dbURL string, batchSize int) {
	defer pool.wg.Done()

	// Each worker gets its own database connection
	db, err := sql.Open("postgres", dbURL)
	if err != nil {
		log.Fatalf("Worker %d: Failed to connect to database: %v", id, err)
	}
	defer db.Close()

	db.SetMaxOpenConns(2)
	db.SetMaxIdleConns(2)

	err = db.Ping()
	if err != nil {
		log.Fatalf("Worker %d: Failed to ping database: %v", id, err)
	}

	log.Printf("Worker %d started", id)

	for batch := range pool.idsChan {
		updated, errors := processBatch(db, batch)

		pool.mu.Lock()
		pool.updateCount += updated
		pool.errorCount += errors
		pool.mu.Unlock()
	}

	log.Printf("Worker %d finished", id)
}

// Process a batch of IDs by updating them in Postgres
func processBatch(db *sql.DB, clipIDs []ClipID) (int64, int64) {
	if len(clipIDs) == 0 {
		return 0, 0
	}

	// Build array of IDs for batch update
	ids := make([]string, len(clipIDs))
	for i, clip := range clipIDs {
		ids[i] = clip.ID
	}

	// Use a transaction for the batch update
	ctx := context.Background()
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		log.Printf("Error starting transaction: %v", err)
		return 0, 1
	}
	defer tx.Rollback()

	// Update query - sets prompt_text to empty string and sets 'prompt' in metadata JSONB to empty string
	// jsonb_set sets the value at the specified path
	query := `UPDATE bots_generatedclip 
			  SET prompt_text = '', 
			      metadata = jsonb_set(metadata, '{prompt}', '""'::jsonb, true)
			  WHERE id = ANY($1)`

	result, err := tx.ExecContext(ctx, query, ids)
	if err != nil {
		log.Printf("Error updating batch: %v", err)
		return 0, 1
	}

	err = tx.Commit()
	if err != nil {
		log.Printf("Error committing transaction: %v", err)
		return 0, 1
	}

	rowsAffected, _ := result.RowsAffected()
	return rowsAffected, 0
}

// Helper function to get environment variable with default value
func getEnv(key, defaultValue string) string {
	value := os.Getenv(key)
	if value == "" {
		return defaultValue
	}
	return value
}

// Helper function to get integer environment variable with default value
func getEnvInt(key string, defaultValue int) int {
	value := os.Getenv(key)
	if value == "" {
		return defaultValue
	}

	var result int
	_, err := fmt.Sscanf(value, "%d", &result)
	if err != nil {
		log.Printf("Warning: Could not parse %s as integer, using default %d", value, defaultValue)
		return defaultValue
	}

	return result
}