module pool

import sync
import sync.stdatomic { new_atomic }
import time

// Eviction channel capacity
const eviction_ch_cap = 1000

// ConnectionPoolable defines the interface for connection objects
pub interface ConnectionPoolable {
mut:
    // validate checks if the connection is still usable
    validate() !bool
    // close terminates the physical connection
    close() !
    // reset returns the connection to initial state for reuse
    reset() !
}

// ConnectionPoolConfig holds configuration settings for the connection pool
@[params]
pub struct ConnectionPoolConfig {
pub mut:
    max_conns          int           = 20               // Maximum allowed connections
    min_idle_conns     int           = 5                // Minimum idle connections to maintain
    max_lifetime       time.Duration = time.hour        // Max lifetime of a connection
    idle_timeout       time.Duration = 30 * time.minute // Time before idle connections are cleaned up
    get_timeout        time.Duration = 5 * time.second  // Max time to wait for a connection
    retry_base_delay   time.Duration = 1 * time.second  // Base delay for retry backoff
    max_retry_delay    time.Duration = 30 * time.second // Maximum delay for retry backoff
    max_retry_attempts int           = 5                // Maximum retry attempts
}

// ConnectionWrapper contains metadata about a pooled connection
struct ConnectionWrapper {
mut:
    conn          &ConnectionPoolable // The actual connection object
    created_at    time.Time           // When connection was created
    last_used_at  time.Time           // Last time connection was used
    last_valid_at time.Time           // Last time connection was validated
    usage_count   int                 // How many times this connection has been used
}

// EvictionPriority indicates urgency of connection cleanup
pub enum EvictionPriority {
    low    // Routine cleanup (connection return)
    medium // Connection get failure
    high   // Configuration change
    urgent // Database/Server recovery
}

// ConnectionPool manages a pool of reusable connections
pub struct ConnectionPool {
mut:
    config ConnectionPoolConfig
    // Lock order:
    // config_mutex > create_mutex > idle_pool_mutex > all_conns_mutex > wait_queue_mutex
    config_mutex     &sync.RwMutex @[required] // Guards configuration changes
    create_mutex     &sync.Mutex // Serializes connection creation
    idle_pool_mutex  &sync.RwMutex              @[required] // Protects idle connections
    all_conns_mutex  &sync.RwMutex              @[required] // Protects all connections map
    wait_queue_mutex &sync.RwMutex              @[required] // Protects wait queue
    is_closed        &stdatomic.AtomicVal[bool] @[required] // Pool shutdown flag
    stop_ch          chan bool             // Signals maintenance thread to stop
    eviction_ch      chan EvictionPriority // Eviction event channel
    cleanup_thread   thread                // Background maintenance thread
    wait_queue       []chan bool           // Client wait queue for connection acquisition
    conn_factory     fn () !&ConnectionPoolable @[required] // Creates new connections
    active_count     &stdatomic.AtomicVal[int]  @[required] // Currently checked-out connections
    created_at       time.Time                      // Pool creation timestamp
    all_conns        map[voidptr]&ConnectionWrapper // All tracked connections
    idle_pool        []&ConnectionWrapper           // Currently idle connections
    creation_errors  &stdatomic.AtomicVal[int] @[required] // Failed creation attempts
    evicted_count    &stdatomic.AtomicVal[int] @[required] // Connections forcibly removed
    creating_count   &stdatomic.AtomicVal[int] @[required] // Connections being created
}

// new_connection_pool creates a new connection pool
pub fn new_connection_pool(conn_factory fn () !&ConnectionPoolable, config ConnectionPoolConfig) !&ConnectionPool {
    // Validate configuration parameters
    check_config(config)!

    mut p := &ConnectionPool{
        conn_factory:     conn_factory
        config:           config
        config_mutex:     sync.new_rwmutex()
        create_mutex:     sync.new_mutex()
        idle_pool_mutex:  sync.new_rwmutex()
        all_conns_mutex:  sync.new_rwmutex()
        wait_queue_mutex: sync.new_rwmutex()
        is_closed:        new_atomic(false)
        stop_ch:          chan bool{cap: 1}
        eviction_ch:      chan EvictionPriority{cap: eviction_ch_cap}
        active_count:     new_atomic(0)
        creation_errors:  new_atomic(0)
        evicted_count:    new_atomic(0)
        creating_count:   new_atomic(0)
        all_conns:        map[voidptr]&ConnectionWrapper{}
    }

    now := time.utc()
    p.created_at = now

    // Initialize minimum idle connections
    for _ in 0 .. config.min_idle_conns {
        conn := p.create_conn_with_retry() or {
            // Cleanup if initialization fails
            p.all_conns_mutex.lock()
            for _, mut wrapper in p.all_conns {
                wrapper.conn.close() or {}
            }
            p.all_conns.clear()
            p.all_conns_mutex.unlock()
            return err
        }
        wrapper := &ConnectionWrapper{
            conn:          conn
            created_at:    now
            last_used_at:  now
            last_valid_at: now
        }
        p.idle_pool << wrapper
        p.all_conns_mutex.lock()
        p.all_conns[conn] = wrapper
        p.all_conns_mutex.unlock()
    }

    // Start background maintenance thread
    p.cleanup_thread = spawn p.background_maintenance()
    // Initial connection pruning
    p.prune_connections()
    return p
}

// create_conn_with_retry creates a connection with exponential backoff retries
fn (mut p ConnectionPool) create_conn_with_retry() !&ConnectionPoolable {
    // Get current configuration
    p.config_mutex.rlock()
    max_attempts := p.config.max_retry_attempts
    base_delay := p.config.retry_base_delay
    max_delay := p.config.max_retry_delay
    p.config_mutex.unlock()

    // Serialize connection creation
    p.create_mutex.lock()
    defer {
        p.create_mutex.unlock()
    }

    mut attempt := 0
    p.creating_count.add(1)
    defer {
        p.creating_count.sub(1)
    }

    for {
        mut conn := p.conn_factory() or {
            // Handle creation error with exponential backoff
            if attempt >= max_attempts {
                return error('Connection creation failed after ${attempt} attempts: ${err}')
            }

            // Calculate next delay with exponential backoff
            mut delay := base_delay * time.Duration(1 << attempt)
            if delay > max_delay {
                delay = max_delay
            }

            time.sleep(delay)
            attempt++
            p.creation_errors.add(1)
            continue
        }

        // Validate new connection
        if !conn.validate() or { false } {
            conn.close() or {}
            return error('New connection validation failed')
        }
        return conn
    }
    return error('Unreachable code')
}

// try_wakeup_waiters attempts to notify waiting clients of available resources
fn (mut p ConnectionPool) try_wakeup_waiters() {
    can_create := p.can_create()
    p.wait_queue_mutex.lock()
    defer {
        p.wait_queue_mutex.unlock()
    }

    // Notify first client if resources are available
    if p.wait_queue.len > 0 {
        if p.idle_pool.len > 0 || can_create {
            to_wake := p.wait_queue[0]
            p.wait_queue.delete(0)
            to_wake <- true
        }
    }
}

// can_create checks if new connections can be created
@[inline]
fn (mut p ConnectionPool) can_create() bool {
    p.config_mutex.rlock()
    max_conns := p.config.max_conns
    p.config_mutex.unlock()
    return p.active_count.load() + p.creating_count.load() < max_conns && !p.is_closed.load()
        && p.all_conns.len < max_conns
}

// get acquires a connection from the pool with timeout
pub fn (mut p ConnectionPool) get() !&ConnectionPoolable {
    start_time := time.utc()
    for {
        // Check if pool is closed
        if p.is_closed.load() {
            return error('Connection pool closed')
        }

        // Try immediate acquisition
        if conn := p.try_get() {
            p.eviction_ch <- .medium
            return conn
        }

        // Check if new connection can be created
        can_create := p.can_create()
        if can_create {
            mut new_conn := p.create_conn_with_retry()!

            // Final check before adding to pool
            if p.is_closed.load() {
                new_conn.close()!
                return error('Connection pool closed')
            }

            p.config_mutex.rlock()
            max_conns := p.config.max_conns
            p.config_mutex.unlock()

            p.all_conns_mutex.lock()
            if p.all_conns.len < max_conns {
                // Successfully create and add new connection
                now := time.utc()
                wrapper := &ConnectionWrapper{
                    conn:          new_conn
                    created_at:    now
                    last_used_at:  now
                    last_valid_at: now
                }
                p.all_conns[new_conn] = wrapper
                p.all_conns_mutex.unlock()
                p.active_count.add(1)
                return new_conn
            } else {
                // Connection limit reached - close new connection
                p.all_conns_mutex.unlock()
                new_conn.close()!
            }
        }

        // Second attempt to get connection
        if conn := p.try_get() {
            return conn
        }

        // Calculate remaining time for connection acquisition
        p.config_mutex.rlock()
        timeout := p.config.get_timeout
        p.config_mutex.unlock()
        elapsed := time.utc() - start_time
        if elapsed > timeout {
            return error('Connection acquisition timeout')
        }
        remaining := timeout - elapsed

        // Set up notification channel
        notify_chan := chan bool{cap: 1}
        defer {
            notify_chan.close()
        }

        // Add to wait queue
        p.wait_queue_mutex.lock()
        p.wait_queue << notify_chan
        p.wait_queue_mutex.unlock()

        select {
            _ := <-notify_chan {
                // Notification received - retry acquisition
            }
            i64(remaining) {
                // Timeout cleanup
                p.wait_queue_mutex.lock()
                for i := 0; i < p.wait_queue.len; i++ {
                    if p.wait_queue[i] == notify_chan {
                        p.wait_queue.delete(i)
                        break
                    }
                }
                p.wait_queue_mutex.unlock()
                if conn := p.try_get() {
                    return conn
                }
                return error('Connection acquisition timeout')
            }
        }
    }

    return error('Unreachable code')
}

// try_get attempts non-blocking connection acquisition
fn (mut p ConnectionPool) try_get() ?&ConnectionPoolable {
    // Get relevant configuration parameters
    p.config_mutex.rlock()
    min_idle := p.config.min_idle_conns
    max_lifetime := p.config.max_lifetime
    p.config_mutex.unlock()

    p.idle_pool_mutex.lock()
    defer {
        p.idle_pool_mutex.unlock()
    }

    // Determine eviction priority based on idle count
    priority := if p.idle_pool.len <= min_idle {
        EvictionPriority.urgent
    } else if p.idle_pool.len > min_idle * 2 {
        EvictionPriority.low
    } else {
        EvictionPriority.medium
    }
    p.eviction_ch <- priority

    // Process idle connections
    for p.idle_pool.len > 0 {
        mut wrapper := p.idle_pool.pop()

        // Check connection lifetime
        age := time.utc() - wrapper.created_at
        if age > max_lifetime {
            // Close expired connection
            p.all_conns_mutex.lock()
            p.all_conns.delete(wrapper.conn)
            p.all_conns_mutex.unlock()
            wrapper.conn.close() or {}
            continue
        }

        // Validate connection
        if !wrapper.conn.validate() or { false } {
            // Handle invalid connection
            p.all_conns_mutex.lock()
            p.all_conns.delete(wrapper.conn)
            p.all_conns_mutex.unlock()
            wrapper.conn.close() or {}
            continue
        }

        wrapper.last_valid_at = time.utc()

        // Mark connection as active
        p.active_count.add(1)
        wrapper.last_used_at = time.utc()
        wrapper.usage_count++
        return wrapper.conn
    }
    return none
}

// put returns a connection to the pool
pub fn (mut p ConnectionPool) put(conn &ConnectionPoolable) ! {
    if p.active_count.load() > 0 {
        // TODO: may need a atomic check here, compare_exchange?
        p.active_count.sub(1)
    }

    mut conn_ptr := unsafe { conn }
    // Handle closed pool case
    if p.is_closed.load() {
        conn_ptr.close()!
        return
    }

    // Reset connection to initial state
    conn_ptr.reset() or {
        conn_ptr.close() or {}
        p.all_conns_mutex.lock()
        p.all_conns.delete(conn)
        p.all_conns_mutex.unlock()
        return err
    }

    p.idle_pool_mutex.lock()
    p.all_conns_mutex.lock()
    defer {
        p.all_conns_mutex.unlock()
        p.idle_pool_mutex.unlock()
    }

    // Return connection to idle pool
    if mut wrapper := p.all_conns[conn] {
        wrapper.last_used_at = time.utc()
        p.idle_pool << wrapper

        // Determine if eviction is needed
        p.config_mutex.rlock()
        low_eviction := p.idle_pool.len > p.config.min_idle_conns
        p.config_mutex.unlock()

        // Wake any waiting clients
        p.try_wakeup_waiters()

        // Trigger eviction if needed
        priority := if low_eviction { EvictionPriority.low } else { EvictionPriority.urgent }
        p.eviction_ch <- priority
    } else {
        // Handle unmanaged connection
        conn_ptr.close()!
        return error('Unmanaged connection')
    }
}

// close shuts down the connection pool and cleans up resources
pub fn (mut p ConnectionPool) close() {
    if p.is_closed.load() {
        return
    }
    p.is_closed.store(true)

    // Signal background thread to stop
    p.stop_ch <- true
    p.cleanup_thread.wait()
    p.stop_ch.close()

    // Close all active connections
    p.idle_pool_mutex.lock()
    p.all_conns_mutex.lock()
    for _, mut wrapper in p.all_conns {
        wrapper.conn.close() or {}
    }
    p.all_conns.clear()
    p.idle_pool.clear()
    p.all_conns_mutex.unlock()
    p.idle_pool_mutex.unlock()

    // Process clients in the wait queue
    p.wait_queue_mutex.lock()
    waiters := p.wait_queue.clone()
    p.wait_queue.clear()
    p.wait_queue_mutex.unlock()

    for ch in waiters {
        ch <- true // Notify all waiting clients
    }

    p.eviction_ch.close()

    // Reset all counters
    p.active_count.store(0)
    p.creation_errors.store(0)
    p.evicted_count.store(0)
    p.creating_count.store(0)
}

// background_maintenance handles periodic connection cleanup
fn (mut p ConnectionPool) background_maintenance() {
    mut first_trigger_time := u64(0)
    mut event_count := 0
    mut min_interval := time.infinite // Dynamic processing interval

    for {
        // Calculate adaptive processing interval
        p.config_mutex.rlock()
        dynamic_interval := if p.config.idle_timeout / 10 > time.second {
            time.second
        } else {
            p.config.idle_timeout / 10
        }
        p.config_mutex.unlock()

        interval := if min_interval < dynamic_interval {
            min_interval
        } else {
            dynamic_interval
        }

        select {
            _ := <-p.stop_ch {
                // Termination signal received
                return
            }
            priority := <-p.eviction_ch {
                // Process event based on priority
                match priority {
                    .low {
                        event_count++
                        min_interval = 100 * time.millisecond
                    }
                    .medium {
                        event_count += 10
                        min_interval = 10 * time.millisecond
                    }
                    .high {
                        event_count += 50
                        min_interval = 1 * time.millisecond
                    }
                    .urgent {
                        event_count += 1000
                        min_interval = 100 * time.microsecond
                    }
                }

                // Track first event time
                if first_trigger_time == 0 {
                    first_trigger_time = time.sys_mono_now()
                }

                elapsed := time.sys_mono_now() - first_trigger_time

                // Determine if immediate processing is needed
                if priority == .urgent
                    || (priority == .high && elapsed > 100 * time.microsecond)
                    || (priority == .medium && elapsed > 1 * time.millisecond)
                    || (priority == .low && elapsed > 10 * time.millisecond)
                    || event_count >= 1000 {
                    p.prune_connections()
                    event_count = 0
                    first_trigger_time = 0
                    min_interval = time.infinite
                }
            }
            i64(interval) {
                // Periodic maintenance
                if event_count > 0 || interval == min_interval {
                    p.prune_connections()
                    event_count = 0
                    first_trigger_time = 0
                    min_interval = time.infinite
                }
            }
        }
    }
}

// prune_connections removes invalid connections and maintains min idle count
fn (mut p ConnectionPool) prune_connections() {
    // Get current configuration parameters
    p.config_mutex.rlock()
    max_lifetime := p.config.max_lifetime
    idle_timeout := p.config.idle_timeout
    min_idle := p.config.min_idle_conns
    p.config_mutex.unlock()

    p.idle_pool_mutex.lock()
    // Remove stale connections
    for i := p.idle_pool.len - 1; i >= 0; i-- {
        mut wrapper := p.idle_pool[i]
        age := time.utc() - wrapper.created_at
        idle_time := time.utc() - wrapper.last_used_at

        if age > max_lifetime || idle_time > idle_timeout || !wrapper.conn.validate() or { false } {
            p.all_conns_mutex.lock()
            p.all_conns.delete(wrapper.conn)
            p.all_conns_mutex.unlock()
            wrapper.conn.close() or {}
            p.idle_pool.delete(i)
            p.evicted_count.add(1)
        } else {
            wrapper.last_valid_at = time.utc()
        }
    }
    current_idle := p.idle_pool.len
    p.idle_pool_mutex.unlock()

    // Calculate connections to create
    to_create := if min_idle > current_idle { min_idle - current_idle } else { 0 }

    // Create needed connections
    mut new_conns := []&ConnectionPoolable{}
    if to_create > 0 {
        for _ in 0 .. to_create {
            if new_conn := p.create_conn_with_retry() {
                new_conns << new_conn
            }
        }
    }

    // Check if pool was closed during creation
    if p.is_closed.load() {
        for mut new_conn in new_conns {
            new_conn.close() or {}
        }
        return
    }

    p.config_mutex.rlock()
    current_min_idle := p.config.min_idle_conns
    max_conns := p.config.max_conns
    p.config_mutex.unlock()

    // Add new connections to the pool
    p.idle_pool_mutex.lock()
    p.all_conns_mutex.lock()
    defer {
        p.all_conns_mutex.unlock()
        p.idle_pool_mutex.unlock()
    }

    actual_needed := if current_min_idle > p.idle_pool.len {
        current_min_idle - p.idle_pool.len
    } else {
        0
    }
    available_slots := max_conns - p.all_conns.len
    mut actual_to_add := if actual_needed > new_conns.len { new_conns.len } else { actual_needed }
    actual_to_add = if actual_to_add > available_slots { available_slots } else { actual_to_add }

    // Create wrapper for each new connection
    for i in 0 .. actual_to_add {
        now := time.utc()
        wrapper := &ConnectionWrapper{
            conn:          new_conns[i]
            created_at:    now
            last_used_at:  now
            last_valid_at: now
        }
        p.idle_pool << wrapper
        p.all_conns[new_conns[i]] = wrapper
    }

    // Close any extra connections
    for i in actual_to_add .. new_conns.len {
        new_conns[i].close() or {}
    }

    // Wake clients if connections were added
    if actual_to_add > 0 {
        p.try_wakeup_waiters()
    }
}

fn check_config(config ConnectionPoolConfig) ! {
    if config.max_conns <= 0 {
        return error('max_conns must be positive')
    }
    if config.min_idle_conns < 0 {
        return error('min_idle_conns cannot be negative')
    }
    if config.min_idle_conns > config.max_conns {
        return error('min_idle_conns cannot exceed max_conns')
    }
    if config.max_lifetime < 0 {
        return error('max_lifetime cannot be negative')
    }
    if config.idle_timeout < 0 {
        return error('idle_timeout cannot be negative')
    }
    if config.idle_timeout > config.max_lifetime {
        return error('idle_timeout cannot exceed max_lifetime')
    }
    if config.get_timeout < 0 {
        return error('get_timeout cannot be negative')
    }
    if config.retry_base_delay < 0 {
        return error('retry_base_delay cannot be negative')
    }
    if config.max_retry_delay < 0 {
        return error('max_retry_delay cannot be negative')
    }
    if config.max_retry_attempts < 0 {
        return error('max_retry_attempts cannot be negative')
    }
}

// update_config changes the connection pool configuration
pub fn (mut p ConnectionPool) update_config(config ConnectionPoolConfig) ! {
    // Validate configuration
    check_config(config)!
    // Check pool status
    if p.is_closed.load() {
        return error('Connection pool is closed')
    }

    // Update configuration
    p.config_mutex.lock()
    p.config = config
    p.config_mutex.unlock()

    // Trigger maintenance
    p.eviction_ch <- .high
}

// signal_recovery_event notifies the pool of recovery event
pub fn (mut p ConnectionPool) signal_recovery_event() {
    p.eviction_ch <- .urgent
}

// send_eviction triggers a cleanup event
pub fn (mut p ConnectionPool) send_eviction(priority EvictionPriority) {
    p.eviction_ch <- priority
}

// ConnectionPoolStats holds statistics about the pool
pub struct ConnectionPoolStats {
pub:
    total_conns     int       // All managed connections
    active_conns    int       // Currently checked-out connections
    idle_conns      int       // Available connections
    waiting_clients int       // Clients waiting for a connection
    evicted_count   int       // Connections forcibly removed
    creation_errors int       // Failed creation attempts
    created_at      time.Time // When pool was created
    creating_count  int       // Connections being created
}

// stats retrieves current connection pool statistics
pub fn (mut p ConnectionPool) stats() ConnectionPoolStats {
    p.idle_pool_mutex.rlock()
    p.all_conns_mutex.rlock()
    p.wait_queue_mutex.rlock()
    defer {
        p.wait_queue_mutex.unlock()
        p.all_conns_mutex.unlock()
        p.idle_pool_mutex.unlock()
    }

    return ConnectionPoolStats{
        total_conns:     p.all_conns.len
        active_conns:    p.active_count.load()
        idle_conns:      p.idle_pool.len
        waiting_clients: p.wait_queue.len
        evicted_count:   p.evicted_count.load()
        creation_errors: p.creation_errors.load()
        created_at:      p.created_at
        creating_count:  p.creating_count.load()
    }
}