package upstream

import (
	"context"
	"fmt"
	cmpp "github.com/bigwhite/gocmpp"
	"time"
)

// A window token belongs to one physical connection and must be released only
// after its Submit attempt completes, preserving per-connection CMPP flow control.

func (p *connectionPool) acquireConnection(ctx context.Context) (*connection, func(), error) {
	waitCtx, cancel := context.WithTimeout(ctx, defaultSubmitTimeout)
	defer cancel()
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()

	for {
		if conn, release := p.tryAcquireConnection(); conn != nil {
			if connected, err := conn.ensureConnected(); err != nil {
				release()
				select {
				case <-waitCtx.Done():
					return nil, nil, waitCtx.Err()
				case <-ticker.C:
					continue
				}
			} else if connected {
				_ = p.reportState(context.Background(), "connected", nil)
			}
			return conn, release, nil
		}
		select {
		case <-waitCtx.Done():
			return nil, nil, waitCtx.Err()
		case <-ticker.C:
		}
	}
}

func (p *connectionPool) tryAcquireConnection() (*connection, func()) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if len(p.conns) == 0 {
		return nil, nil
	}
	for i := 0; i < len(p.conns); i++ {
		index := (p.next + i) % len(p.conns)
		conn := p.conns[index]
		if conn.tryAcquireWindow() {
			p.next = (index + 1) % len(p.conns)
			return conn, conn.releaseWindow
		}
	}
	return nil, nil
}

func (c *connection) tryAcquireWindow() bool {
	if c.window == nil {
		c.window = make(chan struct{}, defaultWindowSize)
	}
	select {
	case c.window <- struct{}{}:
		return true
	default:
		return false
	}
}

func (c *connection) releaseWindow() {
	if c.window == nil {
		return
	}
	select {
	case <-c.window:
	default:
	}
}

func (c *connection) heartbeatLoop(ctx context.Context) {
	interval := time.Duration(c.config.HeartbeatIntervalSeconds) * time.Second
	if interval <= 0 {
		interval = defaultHeartbeatInterval
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if !c.sendHeartbeat() {
				return
			}
		}
	}
}

func (c *connection) sendHeartbeat() bool {
	threshold := c.config.HeartbeatMissThreshold
	if threshold <= 0 {
		threshold = defaultHeartbeatMissThreshold
	}
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return false
	}
	if len(c.heartbeatPending) >= threshold {
		c.mu.Unlock()
		c.handleConnectionLoss(fmt.Errorf("heartbeat timeout after %d unanswered ACTIVE_TEST requests", threshold))
		return false
	}
	if c.client == nil {
		c.mu.Unlock()
		return false
	}
	c.sendMu.Lock()
	seq, err := c.client.SendReqPkt(&cmpp.CmppActiveTestReqPkt{})
	c.sendMu.Unlock()
	if err != nil {
		c.mu.Unlock()
		c.handleConnectionLoss(fmt.Errorf("send ACTIVE_TEST: %w", err))
		return false
	}
	if !c.closed {
		c.heartbeatPending[seq] = time.Now().UTC()
	}
	c.mu.Unlock()
	return true
}

func (c *connection) handleHeartbeatResponse(sequenceID uint32) {
	c.mu.Lock()
	delete(c.heartbeatPending, sequenceID)
	c.mu.Unlock()
}