package health

import (
	"context"
	"errors"
	"sort"
	"strings"
	"sync/atomic"
	"time"

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

var (
	ErrInvalidSchedulerRunner = errors.New("invalid health scheduler runner")
	ErrInvalidDueSource       = errors.New("invalid health due source response")
	ErrInvalidTaskSink        = errors.New("invalid health task sink response")
)

// DueSource is implemented by the Controller's future Redis due-index. Both
// calls are bounded and stay outside the Gateway request path.
type DueSource interface {
	InFlight(context.Context, time.Time) (int, error)
	DueCandidates(context.Context, time.Time, int) ([]Candidate, error)
}

// UpstreamTaskSource scopes due reads and shared capacity to one configured
// upstream while all leases remain in the same Redis namespace.
type UpstreamTaskSource interface {
	InFlightForUpstream(context.Context, string, time.Time) (int, error)
	DueCandidatesForUpstream(context.Context, string, time.Time, int) ([]Candidate, error)
}

// EgressUpstreamTaskSource supplies bounded per-upstream candidates for one
// configured egress probe URL. Its optional nature keeps BASIC-only stores
// source compatible during rollout.
type EgressUpstreamTaskSource interface {
	UpstreamTaskSource
	DueEgressCandidatesForUpstream(context.Context, string, string, time.Time, int) ([]Candidate, error)
}

// TargetUpstreamTaskSource supplies bounded per-upstream candidates for one
// Routing target profile. TARGET facts remain separate from global health.
type TargetUpstreamTaskSource interface {
	UpstreamTaskSource
	DueTargetCandidatesForUpstream(context.Context, string, string, string, time.Time, int) ([]Candidate, error)
}

type upstreamDueSource struct {
	source     UpstreamTaskSource
	upstreamID string
}

func NewUpstreamDueSource(source UpstreamTaskSource, upstreamID string) (DueSource, error) {
	if nilInterface(source) || upstreamID == "" {
		return nil, ErrInvalidSchedulerRunner
	}
	return upstreamDueSource{source: source, upstreamID: upstreamID}, nil
}

func (source upstreamDueSource) InFlight(ctx context.Context, now time.Time) (int, error) {
	return source.source.InFlightForUpstream(ctx, source.upstreamID, now)
}

func (source upstreamDueSource) DueCandidates(ctx context.Context, now time.Time, limit int) ([]Candidate, error) {
	return source.source.DueCandidatesForUpstream(ctx, source.upstreamID, now, limit)
}

type egressUpstreamDueSource struct {
	source     EgressUpstreamTaskSource
	upstreamID string
	targetURL  string
}

func newEgressUpstreamDueSource(source EgressUpstreamTaskSource, upstreamID, targetURL string) (DueSource, error) {
	if nilInterface(source) || upstreamID == "" || targetURL == "" {
		return nil, ErrInvalidSchedulerRunner
	}
	return egressUpstreamDueSource{source: source, upstreamID: upstreamID, targetURL: targetURL}, nil
}

func (source egressUpstreamDueSource) InFlight(ctx context.Context, now time.Time) (int, error) {
	return source.source.InFlightForUpstream(ctx, source.upstreamID, now)
}

func (source egressUpstreamDueSource) DueCandidates(ctx context.Context, now time.Time, limit int) ([]Candidate, error) {
	return source.source.DueEgressCandidatesForUpstream(ctx, source.upstreamID, source.targetURL, now, limit)
}

type targetUpstreamDueSource struct {
	source      TargetUpstreamTaskSource
	upstreamID  string
	routingName string
	targetURL   string
}

func newTargetUpstreamDueSource(
	source TargetUpstreamTaskSource,
	upstreamID string,
	routingName string,
	targetURL string,
) (DueSource, error) {
	profile, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{
		RoutingName: routingName, TargetURL: targetURL,
	})
	if nilInterface(source) || upstreamID == "" || err != nil || profile.RoutingName != routingName || profile.TargetURL != targetURL {
		return nil, ErrInvalidSchedulerRunner
	}
	return targetUpstreamDueSource{
		source: source, upstreamID: upstreamID, routingName: routingName, targetURL: targetURL,
	}, nil
}

func (source targetUpstreamDueSource) InFlight(ctx context.Context, now time.Time) (int, error) {
	return source.source.InFlightForUpstream(ctx, source.upstreamID, now)
}

func (source targetUpstreamDueSource) DueCandidates(ctx context.Context, now time.Time, limit int) ([]Candidate, error) {
	return source.source.DueTargetCandidatesForUpstream(ctx, source.upstreamID, source.routingName, source.targetURL, now, limit)
}

// TaskSink atomically offers an already bounded batch to the shared leased
// task store. It must leave unaccepted candidates eligible for a later tick.
type TaskSink interface {
	Offer(context.Context, []PlannedTask) (int, error)
}

type SchedulerRunnerOptions struct {
	PollInterval time.Duration
	BatchSize    int
	Now          func() time.Time
}

// SchedulerRunner performs one bounded tick at a time. Run has one ticker
// loop, and Tick performs no concurrent fan-out.
type SchedulerRunner struct {
	planner *Planner
	source  DueSource
	sink    TaskSink
	options SchedulerRunnerOptions
}

func NewSchedulerRunner(
	planner *Planner,
	source DueSource,
	sink TaskSink,
	options SchedulerRunnerOptions,
) (*SchedulerRunner, error) {
	if planner == nil || !validSchedulerRunnerParts(source, sink, options) {
		return nil, ErrInvalidSchedulerRunner
	}
	return &SchedulerRunner{planner: planner, source: source, sink: sink, options: options}, nil
}

// ConfiguredSchedulerRunner resolves the current effective upstream check
// policy on every tick. It keeps a running Controller aligned with validated
// config reloads without putting configuration reads in the Gateway path.
type ConfiguredSchedulerRunner struct {
	configuration ConfigurationSource
	upstreamID    string
	source        DueSource
	sink          TaskSink
	options       SchedulerRunnerOptions
}

// ConfiguredSchedulerSupervisor discovers enabled upstreams from each
// published configuration snapshot. It lets newly enabled upstreams begin
// bounded BASIC scheduling without a Controller restart.
type ConfiguredSchedulerSupervisor struct {
	configuration ConfigurationSource
	state         SchedulerStateSource
	source        UpstreamTaskSource
	sink          TaskSink
	options       SchedulerRunnerOptions
	cursor        atomic.Uint64
}

// SchedulerConfigurationSource returns a configuration and its Admin revision
// from one atomic publication. config.Store implements this interface.
type SchedulerConfigurationSource interface {
	ConfigurationSource
	Snapshot() (*config.Config, uint64)
}

// SchedulerStateSource provides the persisted Admin view that fences a
// scheduler from combining two different configuration revisions.
type SchedulerStateSource interface {
	Snapshot(context.Context) (adminstate.Snapshot, error)
}

func NewConfiguredSchedulerRunner(
	configuration ConfigurationSource,
	upstreamID string,
	source DueSource,
	sink TaskSink,
	options SchedulerRunnerOptions,
) (*ConfiguredSchedulerRunner, error) {
	if nilInterface(configuration) || strings.TrimSpace(upstreamID) != upstreamID || upstreamID == "" ||
		!validSchedulerRunnerParts(source, sink, options) {
		return nil, ErrInvalidSchedulerRunner
	}
	return &ConfiguredSchedulerRunner{
		configuration: configuration, upstreamID: upstreamID, source: source, sink: sink, options: options,
	}, nil
}

func NewConfiguredSchedulerSupervisor(
	configuration ConfigurationSource,
	source UpstreamTaskSource,
	sink TaskSink,
	options SchedulerRunnerOptions,
	states ...SchedulerStateSource,
) (*ConfiguredSchedulerSupervisor, error) {
	if nilInterface(configuration) || nilInterface(source) || nilInterface(sink) || options.PollInterval <= 0 ||
		options.BatchSize <= 0 || options.Now == nil || len(states) > 1 || (len(states) == 1 && nilInterface(states[0])) {
		return nil, ErrInvalidSchedulerRunner
	}
	supervisor := &ConfiguredSchedulerSupervisor{
		configuration: configuration, source: source, sink: sink, options: options,
	}
	if len(states) == 1 {
		if _, ok := configuration.(SchedulerConfigurationSource); !ok {
			return nil, ErrInvalidSchedulerRunner
		}
		supervisor.state = states[0]
	}
	return supervisor, nil
}

type TickResult struct {
	Planned int
	Offered int
}

func (runner *SchedulerRunner) Tick(ctx context.Context) (TickResult, error) {
	if ctx == nil || runner == nil || runner.planner == nil || nilInterface(runner.source) ||
		!validSchedulerRunnerParts(runner.source, runner.sink, runner.options) {
		return TickResult{}, ErrInvalidSchedulerRunner
	}
	return tickWithPlanner(ctx, runner.planner, runner.source, runner.sink, runner.options)
}

func (runner *ConfiguredSchedulerRunner) Tick(ctx context.Context) (TickResult, error) {
	if ctx == nil || runner == nil || nilInterface(runner.configuration) || runner.upstreamID == "" ||
		!validSchedulerRunnerParts(runner.source, runner.sink, runner.options) {
		return TickResult{}, ErrInvalidSchedulerRunner
	}
	if err := ctx.Err(); err != nil {
		return TickResult{}, err
	}
	planner, configured, err := configuredPlanner(runner.configuration, runner.upstreamID)
	if err != nil {
		return TickResult{}, err
	}
	if !configured {
		return TickResult{}, nil
	}
	return tickWithPlanner(ctx, planner, runner.source, runner.sink, schedulerOptionsForPlanner(runner.options, planner))
}

func (supervisor *ConfiguredSchedulerSupervisor) Tick(ctx context.Context) (TickResult, error) {
	if ctx == nil || supervisor == nil || nilInterface(supervisor.configuration) || nilInterface(supervisor.source) ||
		nilInterface(supervisor.sink) || supervisor.options.PollInterval <= 0 || supervisor.options.BatchSize <= 0 || supervisor.options.Now == nil {
		return TickResult{}, ErrInvalidSchedulerRunner
	}
	if err := ctx.Err(); err != nil {
		return TickResult{}, err
	}
	configuration, names, routings, err := supervisor.effectiveConfiguration(ctx)
	if err != nil {
		return TickResult{}, err
	}
	result := TickResult{}
	for _, name := range names {
		planner, configured, err := configuredPlannerForConfig(configuration, name)
		if err != nil {
			return TickResult{}, err
		}
		if !configured {
			continue
		}
		basicSource, err := NewUpstreamDueSource(supervisor.source, name)
		if err != nil {
			return TickResult{}, err
		}
		check := config.EffectiveCheck(configuration.Defaults.Check, configuration.Upstreams[name].Check)
		groups, err := supervisor.upstreamDueGroups(basicSource, name, check.URLs, configuredTargetProfiles(configuration, name, routings))
		if err != nil {
			return TickResult{}, err
		}
		remaining := schedulerOptionsForPlanner(supervisor.options, planner).BatchSize
		start := int(supervisor.cursor.Add(1)-1) % len(groups)
		for offset := range groups {
			if remaining <= 0 {
				break
			}
			groupOptions := schedulerOptionsForPlanner(supervisor.options, planner)
			groupOptions.BatchSize = schedulerGroupBatchSize(remaining, len(groups)-offset)
			source := groups[(start+offset)%len(groups)]
			tick, err := tickWithPlanner(ctx, planner, source, supervisor.sink, groupOptions)
			if err != nil {
				return TickResult{}, err
			}
			result.Planned += tick.Planned
			result.Offered += tick.Offered
			remaining -= tick.Offered
		}
	}
	return result, nil
}

// effectiveConfiguration returns the effective Upstream and Routing state for
// one scheduler tick. When Admin state is present, it requires the
// configuration and state snapshots to share a revision; an incomplete or
// torn view schedules nothing.
func (supervisor *ConfiguredSchedulerSupervisor) effectiveConfiguration(
	ctx context.Context,
) (*config.Config, []string, map[string]adminstate.RoutingState, error) {
	if supervisor == nil || nilInterface(supervisor.configuration) {
		return nil, nil, nil, ErrInvalidSchedulerRunner
	}
	if supervisor.state == nil {
		configuration := supervisor.configuration.Current()
		if configuration == nil {
			return nil, nil, nil, ErrInvalidSchedulerRunner
		}
		return configuration, enabledUpstreamNames(configuration), nil, nil
	}
	configurationSource, ok := supervisor.configuration.(SchedulerConfigurationSource)
	if !ok || nilInterface(configurationSource) || nilInterface(supervisor.state) {
		return nil, nil, nil, ErrInvalidSchedulerRunner
	}
	configuration, revision := configurationSource.Snapshot()
	if configuration == nil || revision == 0 {
		return nil, nil, nil, ErrInvalidSchedulerRunner
	}
	snapshot, err := supervisor.state.Snapshot(ctx)
	if err != nil {
		return nil, nil, nil, err
	}
	if snapshot.Config == nil || snapshot.Config.Revision != revision {
		return nil, nil, nil, ErrInvalidSchedulerRunner
	}
	states, err := newSchedulerAdminStates(configuration, snapshot)
	if err != nil {
		return nil, nil, nil, err
	}
	names, err := enabledUpstreamNamesForState(configuration, states.upstreams)
	if err != nil {
		return nil, nil, nil, err
	}
	return configuration, names, states.routings, nil
}

type schedulerAdminStates struct {
	upstreams map[string]adminstate.UpstreamState
	routings  map[string]adminstate.RoutingState
}

func newSchedulerAdminStates(
	configuration *config.Config,
	snapshot adminstate.Snapshot,
) (schedulerAdminStates, error) {
	if configuration == nil {
		return schedulerAdminStates{}, ErrInvalidSchedulerRunner
	}
	states := schedulerAdminStates{
		upstreams: make(map[string]adminstate.UpstreamState, len(snapshot.Upstreams)),
		routings:  make(map[string]adminstate.RoutingState, len(snapshot.Routings)),
	}
	for _, upstream := range snapshot.Upstreams {
		if upstream.Name == "" || upstream.Revision == 0 {
			return schedulerAdminStates{}, ErrInvalidSchedulerRunner
		}
		if _, duplicate := states.upstreams[upstream.Name]; duplicate {
			return schedulerAdminStates{}, ErrInvalidSchedulerRunner
		}
		states.upstreams[upstream.Name] = upstream
	}
	for _, routing := range snapshot.Routings {
		if routing.Name == "" || routing.Revision == 0 {
			return schedulerAdminStates{}, ErrInvalidSchedulerRunner
		}
		if _, duplicate := states.routings[routing.Name]; duplicate {
			return schedulerAdminStates{}, ErrInvalidSchedulerRunner
		}
		states.routings[routing.Name] = routing
	}
	for _, routing := range configuration.Routing {
		if _, exists := states.routings[routing.Name]; !exists {
			return schedulerAdminStates{}, ErrInvalidSchedulerRunner
		}
	}
	return states, nil
}

func (supervisor *ConfiguredSchedulerSupervisor) upstreamDueGroups(
	basic DueSource,
	upstreamID string,
	egressURLs []string,
	targets []configuredTargetProfile,
) ([]DueSource, error) {
	// EGRESS and TARGET due reads lazily materialize references from BASIC due
	// entries. Run them before BASIC so a small batch cannot consume the only
	// source reference before these independent checks are established.
	groups := make([]DueSource, 0, 1+len(egressURLs)+len(targets))
	if egress, supported := supervisor.source.(EgressUpstreamTaskSource); supported {
		for _, targetURL := range egressURLs {
			source, err := newEgressUpstreamDueSource(egress, upstreamID, targetURL)
			if err != nil {
				return nil, err
			}
			groups = append(groups, source)
		}
	}
	if target, supported := supervisor.source.(TargetUpstreamTaskSource); supported {
		for _, profile := range targets {
			source, err := newTargetUpstreamDueSource(target, upstreamID, profile.RoutingName, profile.TargetURL)
			if err != nil {
				return nil, err
			}
			groups = append(groups, source)
		}
	}
	groups = append(groups, basic)
	return groups, nil
}

type configuredTargetProfile struct {
	RoutingName string
	TargetURL   string
}

func configuredTargetProfiles(
	configuration *config.Config,
	upstreamID string,
	routings map[string]adminstate.RoutingState,
) []configuredTargetProfile {
	if configuration == nil || upstreamID == "" {
		return nil
	}
	profiles := make([]configuredTargetProfile, 0)
	for _, route := range configuration.Routing {
		if !route.Enabled || !containsString(route.Upstreams, upstreamID) || (routings != nil && !routings[route.Name].Enabled) {
			continue
		}
		for _, targetURL := range route.Check.Targets {
			profiles = append(profiles, configuredTargetProfile{RoutingName: route.Name, TargetURL: targetURL})
		}
	}
	sort.Slice(profiles, func(left, right int) bool {
		if profiles[left].RoutingName != profiles[right].RoutingName {
			return profiles[left].RoutingName < profiles[right].RoutingName
		}
		return profiles[left].TargetURL < profiles[right].TargetURL
	})
	return profiles
}

func containsString(values []string, target string) bool {
	for _, value := range values {
		if value == target {
			return true
		}
	}
	return false
}

func schedulerGroupBatchSize(batchSize, groups int) int {
	if batchSize <= 0 || groups <= 0 {
		return 0
	}
	result := (batchSize + groups - 1) / groups
	if result < 1 {
		return 1
	}
	return result
}

func tickWithPlanner(
	ctx context.Context,
	planner *Planner,
	source DueSource,
	sink TaskSink,
	options SchedulerRunnerOptions,
) (TickResult, error) {
	if ctx == nil || planner == nil || !validSchedulerRunnerParts(source, sink, options) {
		return TickResult{}, ErrInvalidSchedulerRunner
	}
	if err := ctx.Err(); err != nil {
		return TickResult{}, err
	}
	now := options.Now()
	if now.IsZero() {
		return TickResult{}, ErrInvalidSchedulerRunner
	}
	inFlight, err := source.InFlight(ctx, now.UTC())
	if err != nil {
		return TickResult{}, err
	}
	if inFlight < 0 {
		return TickResult{}, ErrInvalidDueSource
	}
	if inFlight >= planner.policy.MaxInFlight {
		return TickResult{}, nil
	}
	candidates, err := source.DueCandidates(ctx, now.UTC(), options.BatchSize)
	if err != nil {
		return TickResult{}, err
	}
	if len(candidates) > options.BatchSize {
		return TickResult{}, ErrInvalidDueSource
	}
	plans, err := planner.Plan(now.UTC(), inFlight, options.BatchSize, candidates)
	if err != nil {
		return TickResult{}, err
	}
	result := TickResult{Planned: len(plans)}
	if len(plans) == 0 {
		return result, nil
	}
	offered, err := sink.Offer(ctx, plans)
	if err != nil {
		return TickResult{}, err
	}
	if offered < 0 || offered > len(plans) {
		return TickResult{}, ErrInvalidTaskSink
	}
	result.Offered = offered
	return result, nil
}

func (runner *SchedulerRunner) Run(ctx context.Context) error {
	if runner == nil {
		return ErrInvalidSchedulerRunner
	}
	return runScheduler(ctx, runner.options.PollInterval, runner.Tick)
}

func (runner *ConfiguredSchedulerRunner) Run(ctx context.Context) error {
	if runner == nil {
		return ErrInvalidSchedulerRunner
	}
	return runScheduler(ctx, runner.options.PollInterval, runner.Tick)
}

func (supervisor *ConfiguredSchedulerSupervisor) Run(ctx context.Context) error {
	if supervisor == nil {
		return ErrInvalidSchedulerRunner
	}
	return runScheduler(ctx, supervisor.options.PollInterval, supervisor.Tick)
}

func validSchedulerRunnerParts(source DueSource, sink TaskSink, options SchedulerRunnerOptions) bool {
	return !nilInterface(source) && !nilInterface(sink) && options.PollInterval > 0 && options.BatchSize > 0 && options.Now != nil
}

func configuredPlanner(configuration ConfigurationSource, upstreamID string) (*Planner, bool, error) {
	current := configuration.Current()
	if current == nil {
		return nil, false, ErrInvalidSchedulerRunner
	}
	return configuredPlannerForConfig(current, upstreamID)
}

func configuredPlannerForConfig(current *config.Config, upstreamID string) (*Planner, bool, error) {
	if current == nil {
		return nil, false, ErrInvalidSchedulerRunner
	}
	upstream, exists := current.Upstreams[upstreamID]
	if !exists || !upstream.Enabled {
		return nil, false, nil
	}
	check := config.EffectiveCheck(current.Defaults.Check, upstream.Check)
	planner, err := NewPlanner(SchedulePolicy{
		Interval: check.Interval.Value(), Jitter: check.Jitter, MaxInFlight: check.MaxInFlight,
		Timeout: check.Timeout.Value(), MaxAttempts: check.MaxAttempts,
	})
	if err != nil {
		return nil, false, err
	}
	return planner, true, nil
}

func schedulerOptionsForPlanner(options SchedulerRunnerOptions, planner *Planner) SchedulerRunnerOptions {
	if planner.policy.MaxInFlight < options.BatchSize {
		options.BatchSize = planner.policy.MaxInFlight
	}
	return options
}

func enabledUpstreamNames(configuration *config.Config) []string {
	names := make([]string, 0, len(configuration.Upstreams))
	for name, upstream := range configuration.Upstreams {
		if upstream.Enabled {
			names = append(names, name)
		}
	}
	sort.Strings(names)
	return names
}

func enabledUpstreamNamesForState(
	configuration *config.Config,
	states map[string]adminstate.UpstreamState,
) ([]string, error) {
	if configuration == nil {
		return nil, ErrInvalidSchedulerRunner
	}
	names := make([]string, 0, len(configuration.Upstreams))
	for name, upstream := range configuration.Upstreams {
		state, exists := states[name]
		if !exists {
			return nil, ErrInvalidSchedulerRunner
		}
		if upstream.Enabled && state.Enabled {
			names = append(names, name)
		}
	}
	sort.Strings(names)
	return names, nil
}

func runScheduler(ctx context.Context, pollInterval time.Duration, tick func(context.Context) (TickResult, error)) error {
	if ctx == nil || pollInterval <= 0 || tick == nil {
		return ErrInvalidSchedulerRunner
	}
	if _, err := tick(ctx); err != nil {
		return err
	}
	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			if _, err := tick(ctx); err != nil {
				return err
			}
		}
	}
}