package health

import (
	"context"
	"testing"
	"time"

	"proxy-pool/internal/config"
	"proxy-pool/internal/domain/activitypool"
	"proxy-pool/internal/domain/adminstate"
	healthDomain "proxy-pool/internal/domain/health"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

type upstreamDrainPolicySourceStub struct {
	policies []activitypool.UpstreamDrainPolicy
	err      error
}

func (source upstreamDrainPolicySourceStub) ReadUpstreamDrainPolicies(context.Context) ([]activitypool.UpstreamDrainPolicy, error) {
	return append([]activitypool.UpstreamDrainPolicy(nil), source.policies...), source.err
}

type upstreamDrainConfigurationStub struct {
	configuration *config.Config
	revision      uint64
}

type upstreamDrainStateStub struct{ snapshot adminstate.Snapshot }

func (source upstreamDrainStateStub) Snapshot(context.Context) (adminstate.Snapshot, error) {
	return source.snapshot, nil
}

type drainMetricRecord struct {
	reason     healthDomain.DrainReason
	candidates int
	started    int
}

type recordingDrainMetrics struct{ records []drainMetricRecord }

func (metrics *recordingDrainMetrics) ObserveDrain(reason healthDomain.DrainReason, candidates, started int) {
	metrics.records = append(metrics.records, drainMetricRecord{reason: reason, candidates: candidates, started: started})
}

func (source upstreamDrainConfigurationStub) Snapshot() (*config.Config, uint64) {
	return source.configuration, source.revision
}

func TestConfiguredUpstreamDrainReaperStartsBoundedDisabledAssignments(t *testing.T) {
	now := time.Date(2026, 8, 2, 16, 0, 0, 0, time.UTC)
	pool := activitypool.NewMemoryPool()
	if _, err := pool.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1,
		Proxies: []proxyDomain.Proxy{{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080, State: proxyDomain.StateAvailable,
		}},
	}); err != nil {
		t.Fatalf("UpsertFetched() error = %v", err)
	}
	assignment, err := pool.Assign(context.Background(), now, "proxy-a", "worker-a", time.Minute)
	if err != nil {
		t.Fatalf("Assign() error = %v", err)
	}
	policy := activitypool.UpstreamDrainPolicy{UpstreamID: "provider-a", Revision: 21}
	metrics := &recordingDrainMetrics{}
	reaper, err := NewConfiguredUpstreamDrainReaper(upstreamDrainPolicySourceStub{policies: []activitypool.UpstreamDrainPolicy{policy}}, pool,
		UpstreamDrainReaperOptions{PollInterval: time.Second, BatchSize: 4, Now: func() time.Time { return now }, Metrics: metrics})
	if err != nil {
		t.Fatalf("NewConfiguredUpstreamDrainReaper() error = %v", err)
	}
	result, err := reaper.Tick(context.Background())
	if err != nil || result != (UpstreamDrainResult{Candidates: 1, Started: 1}) {
		t.Fatalf("Tick() = %+v, %v", result, err)
	}
	if current, found, err := pool.Get(context.Background(), "proxy-a"); err != nil || !found ||
		current.Epoch != assignment.Epoch || !current.Draining {
		t.Fatalf("Get(after drain) = %+v, %t, %v", current, found, err)
	}
	if len(metrics.records) != 1 || metrics.records[0] != (drainMetricRecord{
		reason: healthDomain.DrainReasonUpstreamDisabled, candidates: 1, started: 1,
	}) {
		t.Fatalf("Drain metrics = %+v", metrics.records)
	}
}

func TestEffectiveUpstreamDrainPolicySourceUsesStableConfigOnlyRevision(t *testing.T) {
	configuration := &config.Config{Upstreams: map[string]config.Upstream{
		"provider-b": {Enabled: true},
		"provider-a": {Enabled: false},
	}}
	source, err := NewEffectiveUpstreamDrainPolicySource(upstreamDrainConfigurationStub{configuration: configuration})
	if err != nil {
		t.Fatalf("NewEffectiveUpstreamDrainPolicySource() error = %v", err)
	}
	policies, err := source.ReadUpstreamDrainPolicies(context.Background())
	if err != nil || len(policies) != 2 || policies[0] != (activitypool.UpstreamDrainPolicy{
		UpstreamID: "provider-a", Revision: 1,
	}) || policies[1] != (activitypool.UpstreamDrainPolicy{
		UpstreamID: "provider-b", Revision: 1, Enabled: true,
	}) {
		t.Fatalf("ReadUpstreamDrainPolicies() = %+v, %v", policies, err)
	}
}

func TestEffectiveUpstreamDrainPolicySourceCombinesConfigurationAndAdminState(t *testing.T) {
	configuration := &config.Config{Upstreams: map[string]config.Upstream{
		"provider-a": {Enabled: false},
		"provider-b": {Enabled: true},
	}}
	source, err := NewEffectiveUpstreamDrainPolicySource(
		upstreamDrainConfigurationStub{configuration: configuration, revision: 9},
		upstreamDrainStateStub{snapshot: adminstate.Snapshot{
			Config: &adminstate.ConfigRevision{Revision: 9},
			Upstreams: []adminstate.UpstreamState{
				{Name: "provider-a", Enabled: true, Revision: 7},
				{Name: "provider-b", Enabled: false, Revision: 8},
			},
		}},
	)
	if err != nil {
		t.Fatalf("NewEffectiveUpstreamDrainPolicySource() error = %v", err)
	}
	policies, err := source.ReadUpstreamDrainPolicies(context.Background())
	want := []activitypool.UpstreamDrainPolicy{
		{UpstreamID: "provider-a", Revision: 7},
		{UpstreamID: "provider-b", Revision: 8},
	}
	if err != nil || len(policies) != len(want) || policies[0] != want[0] || policies[1] != want[1] {
		t.Fatalf("ReadUpstreamDrainPolicies() = %+v, %v; want %+v", policies, err, want)
	}
}