// Package utils provides shared utilities for the fetch_ml project.
package network

import (
	"context"
	"sync"

	"github.com/jfraeys/fetch_ml/internal/logging"
)

type SSHPool struct {
	factory  func() (*SSHClient, error)
	pool     chan *SSHClient
	active   int
	maxConns int
	mu       sync.Mutex
	logger   *logging.Logger
}

func NewSSHPool(maxConns int, factory func() (*SSHClient, error), logger *logging.Logger) *SSHPool {
	return &SSHPool{
		factory:  factory,
		pool:     make(chan *SSHClient, maxConns),
		maxConns: maxConns,
		logger:   logger,
	}
}

func (p *SSHPool) Get(ctx context.Context) (*SSHClient, error) {
	select {
	case conn := <-p.pool:
		return conn, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
		p.mu.Lock()
		if p.active < p.maxConns {
			p.active++
			p.mu.Unlock()
			return p.factory()
		}
		p.mu.Unlock()

		// Wait for available connection
		select {
		case conn := <-p.pool:
			return conn, nil
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}

func (p *SSHPool) Put(conn *SSHClient) {
	select {
	case p.pool <- conn:
	default:
		// Pool is full, close connection
		err := conn.Close()
		if err != nil {
			p.logger.Warn("failed to close SSH connection", "error", err)
		}
		p.mu.Lock()
		p.active--
		p.mu.Unlock()
	}
}

func (p *SSHPool) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()

	// Close all connections in the pool
	close(p.pool)
	for conn := range p.pool {
		err := conn.Close()
		if err != nil {
			p.logger.Warn("failed to close SSH connection", "error", err)
		}
	}

	// Reset active count
	p.active = 0
}