package tlspool

import (
	"math/rand"
	"sync"

	tls_client "github.com/bogdanfinn/tls-client"
	"github.com/bogdanfinn/tls-client/profiles"
)

var defaultProfiles = []profiles.ClientProfile{
	profiles.Chrome_120,
	profiles.Chrome_112,
	profiles.Chrome_110,
	profiles.Safari_IOS_15_5,
	profiles.Chrome_106,
}

// Pool keeps a bank of TLS impersonated HTTP clients per proxy URL ("" means
// direct). Concurrency workers hit different sessions in the bank, and on a
// rotating proxy each session exits via its own IP, which spreads the load
// instead of funnelling every request through one connection. A session is
// replaced after maxUses requests so fingerprints and cookies stay fresh.
type Pool struct {
	mu      sync.Mutex
	banks   map[string]*bank
	size    int
	maxUses int
}

type bank struct {
	sessions []tls_client.HttpClient
	uses     []int
	warmed   []bool
	next     int
}

// New builds a pool with up to size sessions per proxy URL; each session is
// rotated out after maxUses requests. Zero values fall back to 12 and 50.
func New(size, maxUses int) *Pool {
	if size <= 0 {
		size = 12
	}
	if maxUses <= 0 {
		maxUses = 50
	}
	return &Pool{banks: make(map[string]*bank), size: size, maxUses: maxUses}
}

// Get returns the next session for the given proxy, growing the bank lazily
// and rotating individual sessions once their use budget is spent.
func (p *Pool) Get(proxyURL string) tls_client.HttpClient {
	h, _ := p.Take(proxyURL)
	return h
}

// Take behaves like Get and additionally reports needsWarm: true exactly once
// per freshly built session (bank growth, rotation after maxUses, or after a
// Reset). Callers that need session bootstrap cookies send their warmup
// request only when this flag is set, so warmup costs one request per session
// lifetime instead of one per keyword.
func (p *Pool) Take(proxyURL string) (tls_client.HttpClient, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	b := p.banks[proxyURL]
	if b == nil {
		b = &bank{}
		p.banks[proxyURL] = b
	}

	if len(b.sessions) < p.size {
		h := build(proxyURL)
		if h == nil {
			return nil, false
		}
		b.sessions = append(b.sessions, h)
		b.uses = append(b.uses, 1)
		b.warmed = append(b.warmed, true)
		return h, true
	}

	b.next++
	i := b.next % len(b.sessions)
	if b.uses[i] >= p.maxUses {
		if old := b.sessions[i]; old != nil {
			old.CloseIdleConnections()
		}
		if h := build(proxyURL); h != nil {
			b.sessions[i] = h
			b.uses[i] = 0
			b.warmed[i] = false
		}
	}
	b.uses[i]++
	needsWarm := !b.warmed[i]
	b.warmed[i] = true
	return b.sessions[i], needsWarm
}

// Reset drops the whole bank for a proxy so the next Get builds fresh sessions
// with new cookie jars and TLS fingerprints (used after blocks or rate limits).
func (p *Pool) Reset(proxyURL string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if b, ok := p.banks[proxyURL]; ok {
		for _, h := range b.sessions {
			if h != nil {
				h.CloseIdleConnections()
			}
		}
		delete(p.banks, proxyURL)
	}
}

// Size reports the current bank size for a proxy (0 when none built yet).
func (p *Pool) Size(proxyURL string) int {
	p.mu.Lock()
	defer p.mu.Unlock()
	if b, ok := p.banks[proxyURL]; ok {
		return len(b.sessions)
	}
	return 0
}

func build(proxyURL string) tls_client.HttpClient {
	opts := []tls_client.HttpClientOption{
		tls_client.WithTimeoutSeconds(15),
		tls_client.WithClientProfile(defaultProfiles[rand.Intn(len(defaultProfiles))]),
		tls_client.WithCookieJar(tls_client.NewCookieJar()),
		tls_client.WithRandomTLSExtensionOrder(),
	}
	if proxyURL != "" {
		opts = append(opts, tls_client.WithProxyUrl(proxyURL))
	}
	h, err := tls_client.NewHttpClient(tls_client.NewNoopLogger(), opts...)
	if err != nil {
		return nil
	}
	return h
}