package iter

import (
    "runtime"
    "sync/atomic"

    "github.com/sourcegraph/conc"
)

// defaultMaxGoroutines returns the default maximum number of
// goroutines to use within this package.
func defaultMaxGoroutines() int { return runtime.GOMAXPROCS(0) }

// Iterator can be used to configure the behaviour of ForEach
// and ForEachIdx. The zero value is safe to use with reasonable
// defaults.
//
// Iterator is also safe for reuse and concurrent use.
type Iterator[T any] struct {
    // MaxGoroutines controls the maximum number of goroutines
    // to use on this Iterator's methods.
    //
    // If unset, MaxGoroutines defaults to runtime.GOMAXPROCS(0).
    MaxGoroutines int
}

// ForEach executes f in parallel over each element in input.
//
// It is safe to mutate the input parameter, which makes it
// possible to map in place.
//
// ForEach always uses at most runtime.GOMAXPROCS goroutines.
// It takes roughly 2µs to start up the goroutines and adds
// an overhead of roughly 50ns per element of input. For
// a configurable goroutine limit, use a custom Iterator.
func ForEach[T any](input []T, f func(*T)) { Iterator[T]{}.ForEach(input, f) }

// ForEach executes f in parallel over each element in input,
// using up to the Iterator's configured maximum number of
// goroutines.
//
// It is safe to mutate the input parameter, which makes it
// possible to map in place.
//
// It takes roughly 2µs to start up the goroutines and adds
// an overhead of roughly 50ns per element of input.
func (iter Iterator[T]) ForEach(input []T, f func(*T)) {
    iter.ForEachIdx(input, func(_ int, t *T) {
       f(t)
    })
}

// ForEachIdx is the same as ForEach except it also provides the
// index of the element to the callback.
func ForEachIdx[T any](input []T, f func(int, *T)) { Iterator[T]{}.ForEachIdx(input, f) }

// ForEachIdx is the same as ForEach except it also provides the
// index of the element to the callback.
func (iter Iterator[T]) ForEachIdx(input []T, f func(int, *T)) {
    if iter.MaxGoroutines == 0 {
       // iter is a value receiver and is hence safe to mutate
       iter.MaxGoroutines = defaultMaxGoroutines()
    }

    numInput := len(input)
    if iter.MaxGoroutines > numInput {
       // No more concurrent tasks than the number of input items.
       iter.MaxGoroutines = numInput
    }

    var idx atomic.Int64
    // Create the task outside the loop to avoid extra closure allocations.
    task := func() {
       i := int(idx.Add(1) - 1)
       for ; i < numInput; i = int(idx.Add(1) - 1) {
          f(i, &input[i])
       }
    }

    var wg conc.WaitGroup
    for i := 0; i < iter.MaxGoroutines; i++ {
       wg.Go(task)
    }
    wg.Wait()
}

func main() {
    // Let's imagine we need to send emails to multiple customers
    customers := []struct {
        ID    int
        Email string
        Name  string
        Type  string
    }{
        {1, "customer1@example.com", "Alex Johnson", "Premium"},
        {2, "customer2@example.com", "Sam Smith", "Regular"},
        {3, "customer3@example.com", "Taylor Brown", "Premium"},
        {4, "customer4@example.com", "Jordan Lee", "Trial"},
        {5, "customer5@example.com", "Casey Miller", "Regular"},
    }
    
    var wg sync.WaitGroup
    
    // Launch a goroutine for each email that needs to be sent
    for _, customer := range customers {
        wg.Add(1)
        // Simply launch the goroutine with our function
        go sendCustomerEmail(customer, &wg)
    }
    
    // Wait for all emails to be sent
    wg.Wait()
    fmt.Println("All emails have been sent!")
}

// A regular function that sends an email to a customer
func sendCustomerEmail(c struct {
    ID    int
    Email string
    Name  string
    Type  string
}, wg *sync.WaitGroup) {
    defer wg.Done()
    
    // Choose email template based on customer type
    var template string
    switch c.Type {
    case "Premium":
        template = "premium_newsletter.html"
    case "Trial":
        template = "trial_extension_offer.html"
    default:
        template = "regular_newsletter.html"
    }
    
    // Log the start of sending
    fmt.Printf("Sending email to %s (%s)\n", c.Name, c.Email)
    
    // Simulate the actual email sending
    time.Sleep(time.Millisecond * 300) // Network/API delay
    
    // In a real app, we might have code like:
    // emailContent := generateEmailContent(template, c.Name)
    // response, err := emailClient.Send(c.Email, "March Newsletter", emailContent)
    
    // Log the completion
    fmt.Printf("✓ Email sent to %s successfully\n", c.Name)
}