module pg

import sync
import time

// PoolStats reports the current state of a `DB`'s connection pool.
pub struct PoolStats {
pub:
    max_open_connections int // configured limit (0 = unlimited)
    open_connections     int // total in-use + idle conns
    in_use               int // conns currently checked out
    idle                 int // conns parked as idle
    wait_count           int // number of callers currently blocked on acquire
}

// PoolConfig configures pool behavior at construction time.
@[params]
pub struct PoolConfig {
pub:
    max_open_conns    int // 0 = unlimited
    max_idle_conns    int = 2 // 0 = keep no idle conns
    conn_max_lifetime time.Duration // 0 = unlimited
}

// IdleSlot is the pool's internal representation of a pooled libpq handle.
// The pool tracks raw PGconn* metadata in idle/waiter channels; `acquire`
// allocates a fresh `&Conn` wrapping the slot and `release` extracts the
// metadata back. Keeping the wrapper separate from the slot means a stale
// `&Conn` kept by user code after `close()` cannot be revived even when
// the pool re-hands the same physical PGconn* to another caller.
struct IdleSlot {
    handle     voidptr
    created_at time.Time
mut:
    bad bool
}

@[heap]
struct Pool {
mut:
    mu           &sync.Mutex = unsafe { nil }
    conninfo     string
    max_open     int
    max_idle     int
    max_lifetime time.Duration
    idle         []IdleSlot
    open_count   int
    waiters      []chan IdleSlot
    closed       bool
    // last_ids stores the per-thread last-inserted id so DB.last_id() returns
    // the value captured on the same pooled conn that ran the INSERT, instead
    // of calling LASTVAL() on whatever conn the pool happens to hand out next
    // (which is session-scoped and would return 0 or a wrong value).
    last_ids_mu &sync.Mutex = unsafe { nil }
    last_ids    map[u64]i64
}

fn new_pool(conninfo string, cfg PoolConfig) &Pool {
    mut p := &Pool{
        mu:           sync.new_mutex()
        last_ids_mu:  sync.new_mutex()
        conninfo:     conninfo
        max_open:     cfg.max_open_conns
        max_idle:     cfg.max_idle_conns
        max_lifetime: cfg.conn_max_lifetime
    }
    if p.max_open < 0 {
        p.max_open = 0
    }
    if p.max_idle < 0 {
        p.max_idle = 0
    }
    return p
}

// wrap_slot builds the fresh `&Conn` that gets handed to a caller. Each
// acquire allocates a new wrapper so a previously-checked-out (and since
// released) `&Conn` reference can never be reused to reach the underlying
// PGconn*.
fn (p &Pool) wrap_slot(slot IdleSlot) &Conn {
    return &Conn{
        conn:       slot.handle
        pool:       unsafe { p }
        created_at: slot.created_at
        bad:        slot.bad
    }
}

fn (mut p Pool) acquire() !&Conn {
    for {
        p.mu.lock()
        if p.closed {
            p.mu.unlock()
            return error('pg: pool is closed')
        }
        // Pop newest idle slot (LIFO keeps the freshest connection warm)
        for p.idle.len > 0 {
            slot := p.idle.last()
            p.idle.delete_last()
            if slot_expired(slot, p.max_lifetime) || slot_bad(slot) {
                physical_close_handle(slot.handle)
                p.open_count--
                continue
            }
            p.mu.unlock()
            return p.wrap_slot(slot)
        }
        // Capacity available: open a new conn outside the lock
        if p.max_open == 0 || p.open_count < p.max_open {
            p.open_count++
            conninfo := p.conninfo
            p.mu.unlock()
            slot := connect_slot(conninfo) or {
                p.mu.lock()
                p.open_count--
                p.mu.unlock()
                return err
            }
            // close() may have run while we were dialing outside the lock.
            // Honor the close contract by tearing the fresh handle down
            // instead of returning a live Conn after shutdown.
            p.mu.lock()
            if p.closed {
                p.open_count--
                p.mu.unlock()
                physical_close_handle(slot.handle)
                return error('pg: pool is closed')
            }
            p.mu.unlock()
            return p.wrap_slot(slot)
        }
        // At capacity: wait for a release/close/cap-raise to signal us.
        // Senders transfer slot ownership via this channel; a sentinel slot
        // with a nil handle means "capacity changed, retry the acquire
        // loop" (used by set_max_open when raising the limit).
        waiter := chan IdleSlot{cap: 1}
        p.waiters << waiter
        p.mu.unlock()
        slot := <-waiter or { return error('pg: pool was closed while waiting for connection') }
        if isnil(slot.handle) {
            continue
        }
        return p.wrap_slot(slot)
    }
    return error('pg: unreachable')
}

fn (mut p Pool) release(conn &Conn) {
    if isnil(conn) {
        return
    }
    mut c := unsafe { conn }
    // Detach the wrapper from the physical handle before doing anything else.
    // This is what makes a stale `&Conn` reference held by user code inert:
    // any subsequent method call on it will see `c.conn == nil` and error.
    // It is also the idempotency guard against double-close — a second
    // release() finds nothing to return to the pool and is a no-op.
    if isnil(c.conn) {
        return
    }
    slot := IdleSlot{
        handle:     c.conn
        created_at: c.created_at
        bad:        c.bad
    }
    c.conn = unsafe { nil }
    c.pool = unsafe { nil }
    p.mu.lock()
    if p.closed {
        // close() only decremented open_count for slots it found parked;
        // in-use conns are accounted for here as they trickle back.
        p.open_count--
        p.mu.unlock()
        physical_close_handle(slot.handle)
        return
    }
    // Discard broken or expired conns
    if slot_bad(slot) || slot_expired(slot, p.max_lifetime) {
        physical_close_handle(slot.handle)
        p.open_count--
        // A capacity slot just opened; hand it to a waiter as a fresh conn
        if p.waiters.len > 0 && (p.max_open == 0 || p.open_count < p.max_open) {
            waiter := p.waiters[0]
            p.waiters.delete(0)
            p.open_count++
            conninfo := p.conninfo
            p.mu.unlock()
            new_slot := connect_slot(conninfo) or {
                p.mu.lock()
                p.open_count--
                // Capacity is now open but the dial just failed. Wake every
                // other parked waiter too, otherwise they hang forever waiting
                // for a release that will never come (e.g. max_open=1).
                extras := p.waiters.clone()
                p.waiters = []chan IdleSlot{}
                p.mu.unlock()
                waiter.close()
                for w in extras {
                    w.close()
                }
                return
            }
            p.mu.lock()
            if p.closed {
                // Pool closed during the dial: drop the new conn and signal the waiter.
                p.open_count--
                p.mu.unlock()
                physical_close_handle(new_slot.handle)
                waiter.close()
                return
            }
            waiter <- new_slot
            p.mu.unlock()
            return
        }
        p.mu.unlock()
        return
    }
    // If a recent set_max_open() shrank the cap below open_count, this
    // returning slot has to be retired even when a waiter is queued —
    // otherwise the fast hand-off keeps open_count pinned above the new
    // limit forever under steady traffic. Wake the waiter with a retry
    // sentinel so it re-evaluates capacity under the new cap.
    if p.max_open > 0 && p.open_count > p.max_open {
        physical_close_handle(slot.handle)
        p.open_count--
        if p.waiters.len > 0 {
            waiter := p.waiters[0]
            p.waiters.delete(0)
            waiter <- IdleSlot{
                handle: unsafe { nil }
            }
        }
        p.mu.unlock()
        return
    }
    // Healthy conn: prefer handing it directly to a waiter. Send under the
    // lock (cap:1 makes it non-blocking) so close() can't slip in and orphan
    // the popped waiter with a live conn.
    if p.waiters.len > 0 {
        waiter := p.waiters[0]
        p.waiters.delete(0)
        waiter <- slot
        p.mu.unlock()
        return
    }
    // Park as idle, unless we'd exceed max_idle (0 = keep no idle conns)
    // or we are over a recently-shrunk max_open and need to converge.
    if p.idle.len >= p.max_idle || (p.max_open > 0 && p.open_count > p.max_open) {
        physical_close_handle(slot.handle)
        p.open_count--
        p.mu.unlock()
        return
    }
    p.idle << slot
    p.mu.unlock()
}

fn (mut p Pool) close() {
    p.mu.lock()
    if p.closed {
        p.mu.unlock()
        return
    }
    p.closed = true
    for slot in p.idle {
        physical_close_handle(slot.handle)
    }
    idle_len := p.idle.len
    p.idle = []IdleSlot{}
    p.open_count -= idle_len
    for waiter in p.waiters {
        waiter.close()
    }
    p.waiters = []chan IdleSlot{}
    p.mu.unlock()
}

fn (mut p Pool) stats() PoolStats {
    p.mu.lock()
    stats := PoolStats{
        max_open_connections: p.max_open
        open_connections:     p.open_count
        idle:                 p.idle.len
        in_use:               p.open_count - p.idle.len
        wait_count:           p.waiters.len
    }
    p.mu.unlock()
    return stats
}

fn (mut p Pool) set_max_open(n int) {
    mut nn := n
    if nn < 0 {
        nn = 0
    }
    p.mu.lock()
    p.max_open = nn
    // Shrink: drop idle slots until we are back under the new cap. Without
    // this, acquire() pops idle slots before checking max_open, so two
    // callers could still both succeed after set_max_open_conns(1) when two
    // warm conns are parked. In-use conns can't be reclaimed here — they get
    // discarded on release once open_count is over the cap.
    if nn > 0 {
        for p.idle.len > 0 && p.open_count > nn {
            slot := p.idle.last()
            p.idle.delete_last()
            physical_close_handle(slot.handle)
            p.open_count--
        }
    }
    // Raise: wake parked waiters so they can retry the acquire loop and
    // dial against the new headroom. Without this nudge they stay blocked
    // until some other release happens, which may never come if all current
    // conns are long-lived.
    if p.waiters.len > 0 && (nn == 0 || p.open_count < nn) {
        spare := if nn == 0 { p.waiters.len } else { nn - p.open_count }
        mut n_wake := if spare < p.waiters.len { spare } else { p.waiters.len }
        for n_wake > 0 {
            waiter := p.waiters[0]
            p.waiters.delete(0)
            waiter <- IdleSlot{
                handle: unsafe { nil }
            }
            n_wake--
        }
    }
    p.mu.unlock()
}

fn (mut p Pool) set_max_idle(n int) {
    mut nn := n
    if nn < 0 {
        nn = 0
    }
    p.mu.lock()
    p.max_idle = nn
    for p.idle.len > nn {
        slot := p.idle.last()
        p.idle.delete_last()
        physical_close_handle(slot.handle)
        p.open_count--
    }
    p.mu.unlock()
}

fn (mut p Pool) set_conn_max_lifetime(d time.Duration) {
    p.mu.lock()
    p.max_lifetime = d
    p.mu.unlock()
}

// stash_last_id records `id` as the last-inserted id for the current thread.
// Called by `DB.insert` right after running INSERT on the pinned conn, so the
// next `DB.last_id()` call on the same thread returns this exact value rather
// than running LASTVAL() on a different pooled session.
fn (mut p Pool) stash_last_id(id int) {
    tid := sync.thread_id()
    p.last_ids_mu.lock()
    p.last_ids[tid] = i64(id)
    p.last_ids_mu.unlock()
}

// take_last_id returns the last-inserted id stashed by this thread's most
// recent `DB.insert`, or 0 if there is none. Consuming on read keeps the
// map bounded as threads come and go, and prevents a future thread that
// reuses this tid from observing a stale id before doing its own insert.
fn (mut p Pool) take_last_id() int {
    tid := sync.thread_id()
    p.last_ids_mu.lock()
    id := p.last_ids[tid] or { i64(0) }
    p.last_ids.delete(tid)
    p.last_ids_mu.unlock()
    return int(id)
}