package store

import (
	"context"
	"errors"
	"sync"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

const leadershipReleaseTimeout = 5 * time.Second

const (
	identityCoordinatorLeadershipKey    = "easyai-gateway-identity-coordinator"
	securityEventHeartbeatLeadershipKey = "easyai-gateway-security-event-heartbeat"
)

// Leadership holds a PostgreSQL session advisory lock. Callers must keep the
// lease alive and release it when leadership is no longer needed.
type Leadership interface {
	KeepAlive(context.Context) error
	Release()
}

type postgresLeadership struct {
	conn *pgxpool.Conn
	key  string
	once sync.Once
}

func (s *Store) TryAcquireIdentityCoordinatorLeadership(ctx context.Context) (Leadership, bool, error) {
	return s.tryAcquireLeadership(ctx, identityCoordinatorLeadershipKey)
}

func (s *Store) TryAcquireSecurityEventHeartbeatLeadership(ctx context.Context) (Leadership, bool, error) {
	return s.tryAcquireLeadership(ctx, securityEventHeartbeatLeadershipKey)
}

func (s *Store) tryAcquireLeadership(ctx context.Context, key string) (Leadership, bool, error) {
	conn, err := s.pool.Acquire(ctx)
	if err != nil {
		return nil, false, err
	}
	var acquired bool
	if err := conn.QueryRow(ctx,
		`SELECT pg_try_advisory_lock(hashtextextended($1, 0))`,
		key,
	).Scan(&acquired); err != nil {
		conn.Release()
		return nil, false, err
	}
	if !acquired {
		conn.Release()
		return nil, false, nil
	}
	return &postgresLeadership{conn: conn, key: key}, true, nil
}

func (leadership *postgresLeadership) KeepAlive(ctx context.Context) error {
	if leadership == nil || leadership.conn == nil {
		return errors.New("leadership lease is closed")
	}
	var alive int
	return leadership.conn.QueryRow(ctx, `SELECT 1`).Scan(&alive)
}

func (leadership *postgresLeadership) Release() {
	if leadership == nil {
		return
	}
	leadership.once.Do(func() {
		unlockCtx, unlockCancel := context.WithTimeout(context.Background(), leadershipReleaseTimeout)
		var unlocked bool
		err := leadership.conn.QueryRow(unlockCtx,
			`SELECT pg_advisory_unlock(hashtextextended($1, 0))`,
			leadership.key,
		).Scan(&unlocked)
		unlockCancel()
		if err != nil || !unlocked {
			// A session-scoped advisory lock must never return to the pool when
			// unlock is uncertain. Closing the connection lets PostgreSQL
			// release the lock and forces the pool to replace the session.
			conn := leadership.conn.Hijack()
			closeCtx, closeCancel := context.WithTimeout(context.Background(), leadershipReleaseTimeout)
			_ = conn.Close(closeCtx)
			closeCancel()
		} else {
			leadership.conn.Release()
		}
		leadership.conn = nil
	})
}