// Package sched schedules (target, check-family) jobs with a single priority timer wheel.
//
// Guarantees (docs/PROBE-PROTOCOL.md "Rules the agent follows" + SPEC ethics):
//   - per-family intervals: http → schedule.tiers[tier], dns → dns_every, ping/tcp → ping_every,
//     traceroute → traceroute_every; boost multiplies the interval for listed targets until `until`;
//   - ±10 % jitter on every interval, initial offsets spread uniformly over the interval (no start-up burst);
//   - never below MinInterval (10 s) even if the server asks for less (clamped, logged);
//   - a global concurrency limit, at most one in-flight check per hostname, at most one traceroute at a time;
//   - config changes (targets added/removed/edited, intervals) are applied live without restarting.
//
// All state is owned by the Run loop goroutine; Apply/SetBoost communicate through channels, so no locks.
package sched

import (
	"container/heap"
	"context"
	"log/slog"
	"math/rand/v2"
	"runtime/debug"
	"strconv"
	"sync/atomic"
	"time"

	"internetpressure.io/probe-agent/internal/protocol"
)

// Family is a check family.
type Family string

// Families.
const (
	FamilyHTTP       Family = "http"
	FamilyDNS        Family = "dns"
	FamilyPing       Family = "ping"
	FamilyTCP        Family = "tcp"
	FamilyTraceroute Family = "traceroute"
)

// Defaults when the server omits a schedule value.
const (
	MinInterval          = 10 * time.Second
	DefaultDNSEvery      = 60 * time.Second
	DefaultPingEvery     = 30 * time.Second
	DefaultTraceEvery    = 15 * time.Minute
	DefaultTierInterval  = 180 * time.Second
	JitterFraction       = 0.10
	busyRetryBase        = 500 * time.Millisecond
	DefaultMaxConcurrent = 8
)

// Job is one schedulable unit.
type Job struct {
	Key      string // target_id + "/" + family
	Target   protocol.Target
	Family   Family
	Interval time.Duration // base interval (clamped, before boost/jitter)
}

// Runner executes a job. It must return when ctx is done.
type Runner func(ctx context.Context, job Job)

// Options tune the scheduler.
type Options struct {
	MaxConcurrency int
	MinInterval    time.Duration // default MinInterval; tests lower it
	Logger         *slog.Logger
	Seed           uint64 // 0 → random
	// Capabilities: when false the family is not scheduled at all.
	CanPing       bool
	CanTraceroute bool
	// IntervalScale multiplies every configured interval before the clamp (tests only; 0 → 1).
	IntervalScale float64
}

type entry struct {
	job     Job
	gen     uint64
	running bool
}

type item struct {
	key string
	gen uint64
	at  time.Time
}

type itemHeap []item

func (h itemHeap) Len() int           { return len(h) }
func (h itemHeap) Less(i, j int) bool { return h[i].at.Before(h[j].at) }
func (h itemHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *itemHeap) Push(x any)        { *h = append(*h, x.(item)) }
func (h *itemHeap) Pop() any          { old := *h; n := len(old); it := old[n-1]; *h = old[:n-1]; return it }
func (h itemHeap) peek() (item, bool) {
	if len(h) == 0 {
		return item{}, false
	}
	return h[0], true
}

type doneMsg struct {
	key      string
	gen      uint64
	host     string
	trace    bool
	panicked bool
}

// Scheduler is the timer wheel.
type Scheduler struct {
	run  Runner
	opts Options
	rng  *rand.Rand
	log  *slog.Logger

	applyCh chan *protocol.RemoteConfig
	boostCh chan *protocol.Boost
	doneCh  chan doneMsg

	// loop-owned
	entries  map[string]*entry
	heap     itemHeap
	hostBusy map[string]int
	inflight int
	trBusy   bool
	boost    *protocol.Boost
	gen      uint64

	targets atomic.Int32
	jobs    atomic.Int32
	panics  atomic.Uint64
	runs    atomic.Uint64
}

// New creates a scheduler; call Run to start it.
func New(run Runner, opts Options) *Scheduler {
	if opts.MaxConcurrency <= 0 {
		opts.MaxConcurrency = DefaultMaxConcurrent
	}
	if opts.MinInterval <= 0 {
		opts.MinInterval = MinInterval
	}
	if opts.Logger == nil {
		opts.Logger = slog.Default()
	}
	seed := opts.Seed
	if seed == 0 {
		seed = rand.Uint64()
	}
	return &Scheduler{
		run:      run,
		opts:     opts,
		rng:      rand.New(rand.NewPCG(seed, seed^0x9e3779b97f4a7c15)),
		log:      opts.Logger,
		applyCh:  make(chan *protocol.RemoteConfig, 1),
		boostCh:  make(chan *protocol.Boost, 1),
		doneCh:   make(chan doneMsg, 64),
		entries:  map[string]*entry{},
		hostBusy: map[string]int{},
	}
}

// Apply installs a new remote config (targets + schedule + boost). Non-blocking; the latest wins.
func (s *Scheduler) Apply(cfg *protocol.RemoteConfig) {
	for {
		select {
		case s.applyCh <- cfg:
			return
		default:
			select {
			case <-s.applyCh: // drop the stale pending config
			default:
			}
		}
	}
}

// SetBoost replaces the active boost (nil clears it).
func (s *Scheduler) SetBoost(b *protocol.Boost) {
	select {
	case s.boostCh <- b:
	default:
		select {
		case <-s.boostCh:
		default:
		}
		s.boostCh <- b
	}
}

// Targets returns the number of distinct scheduled targets.
func (s *Scheduler) Targets() int { return int(s.targets.Load()) }

// Jobs returns the number of scheduled (target, family) jobs.
func (s *Scheduler) Jobs() int { return int(s.jobs.Load()) }

// Panics returns the number of recovered check panics.
func (s *Scheduler) Panics() uint64 { return s.panics.Load() }

// Runs returns the number of completed job executions.
func (s *Scheduler) Runs() uint64 { return s.runs.Load() }

// Run drives the wheel until ctx is cancelled. In-flight jobs are cancelled through ctx too.
func (s *Scheduler) Run(ctx context.Context) {
	timer := time.NewTimer(time.Hour)
	timer.Stop()
	for {
		// Dispatch everything that is due, then arm the timer for the next item.
		now := time.Now()
		for {
			it, ok := s.heap.peek()
			if !ok || it.at.After(now) {
				break
			}
			heap.Pop(&s.heap)
			s.dispatch(ctx, it, now)
		}
		if it, ok := s.heap.peek(); ok {
			timer.Reset(time.Until(it.at))
		} else {
			timer.Stop()
		}
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case cfg := <-s.applyCh:
			s.apply(cfg, time.Now())
		case b := <-s.boostCh:
			s.boost = b
		case d := <-s.doneCh:
			s.complete(d, time.Now())
		case <-timer.C:
		}
	}
}

func (s *Scheduler) dispatch(ctx context.Context, it item, now time.Time) {
	e, ok := s.entries[it.key]
	if !ok || e.gen != it.gen {
		return // removed or replaced
	}
	host := e.job.Target.Hostname
	trace := e.job.Family == FamilyTraceroute
	if e.running || s.hostBusy[host] > 0 || s.inflight >= s.opts.MaxConcurrency || (trace && s.trBusy) {
		// Busy: try again shortly, spread so that many deferred jobs don't collide. The retry window is
		// 0.5–1.5 s in production (MinInterval 10 s) and scales down with MinInterval in tests.
		base := busyRetryBase
		if b := s.opts.MinInterval / 20; b < base {
			base = b
		}
		s.push(it.key, it.gen, now.Add(base+time.Duration(s.rng.Int64N(int64(2*base)))))
		return
	}
	e.running = true
	s.hostBusy[host]++
	s.inflight++
	if trace {
		s.trBusy = true
	}
	job := e.job
	gen := e.gen
	go func() {
		msg := doneMsg{key: job.Key, gen: gen, host: host, trace: trace}
		defer func() {
			if r := recover(); r != nil {
				msg.panicked = true
				s.panics.Add(1)
				s.log.Error("check panicked", "job", job.Key, "panic", r, "stack", string(debug.Stack()))
			}
			s.runs.Add(1)
			select {
			case s.doneCh <- msg:
			case <-ctx.Done():
				// Loop has exited; nothing to release.
			}
		}()
		s.run(ctx, job)
	}()
}

func (s *Scheduler) complete(d doneMsg, now time.Time) {
	s.inflight--
	if s.hostBusy[d.host] > 1 {
		s.hostBusy[d.host]--
	} else {
		delete(s.hostBusy, d.host)
	}
	if d.trace {
		s.trBusy = false
	}
	e, ok := s.entries[d.key]
	if !ok {
		return // removed while running
	}
	if e.gen != d.gen {
		return // replaced while running: the new generation already has a heap item
	}
	e.running = false
	s.push(d.key, e.gen, now.Add(s.nextDelay(e.job, now)))
}

func (s *Scheduler) push(key string, gen uint64, at time.Time) {
	heap.Push(&s.heap, item{key: key, gen: gen, at: at})
}

// nextDelay = interval × boost (if active) clamped to MinInterval, with ±10 % jitter.
func (s *Scheduler) nextDelay(job Job, now time.Time) time.Duration {
	base := job.Interval
	if s.boost != nil && s.boost.Active(job.Target.TargetID, now) {
		base = time.Duration(float64(base) * s.boost.Factor)
	}
	if base < s.opts.MinInterval {
		base = s.opts.MinInterval
	}
	return Jitter(base, s.rng)
}

// Jitter returns d × U(0.9, 1.1).
func Jitter(d time.Duration, rng *rand.Rand) time.Duration {
	f := 1 + (rng.Float64()*2-1)*JitterFraction
	return time.Duration(float64(d) * f)
}

// InitialOffset returns U[0, d): uniform spread of first runs over one interval.
func InitialOffset(d time.Duration, rng *rand.Rand) time.Duration {
	if d <= 0 {
		return 0
	}
	return time.Duration(rng.Int64N(int64(d)))
}

// apply computes the desired job set from cfg and reconciles it with the current entries.
func (s *Scheduler) apply(cfg *protocol.RemoteConfig, now time.Time) {
	desired := map[string]Job{}
	targets := 0
	if cfg != nil && cfg.Probe.Enabled {
		for _, j := range JobsFor(cfg, s.opts) {
			desired[j.Key] = j
		}
		targets = len(cfg.Targets)
	}
	// Remove vanished jobs (heap items are dropped lazily on pop).
	for key := range s.entries {
		if _, keep := desired[key]; !keep {
			delete(s.entries, key)
		}
	}
	added, changed := 0, 0
	for key, job := range desired {
		e, ok := s.entries[key]
		if !ok {
			s.gen++
			s.entries[key] = &entry{job: job, gen: s.gen}
			s.push(key, s.gen, now.Add(InitialOffset(job.Interval, s.rng)))
			added++
			continue
		}
		if targetChanged(e.job.Target, job.Target) {
			// Re-key the generation: the old heap item is discarded on pop, the job restarts spread out.
			s.gen++
			e.job, e.gen = job, s.gen
			s.push(key, s.gen, now.Add(InitialOffset(job.Interval, s.rng)))
			changed++
			continue
		}
		e.job.Interval = job.Interval // takes effect after the next run
	}
	s.targets.Store(int32(targets))
	s.jobs.Store(int32(len(s.entries)))
	if cfg != nil {
		s.log.Info("schedule applied", "config_version", cfg.ConfigVersion, "enabled", cfg.Probe.Enabled,
			"targets", targets, "jobs", len(s.entries), "added", added, "changed", changed)
	}
}

func targetChanged(a, b protocol.Target) bool {
	if a.Hostname != b.Hostname || a.URL != b.URL || a.Port != b.Port || a.FixedIP() != b.FixedIP() || a.Tier != b.Tier {
		return true
	}
	return false
}

// JobsFor expands a remote config into jobs (one per target × family), applying defaults, capability filters
// and the MinInterval clamp.
func JobsFor(cfg *protocol.RemoteConfig, opts Options) []Job {
	minI := opts.MinInterval
	if minI <= 0 {
		minI = MinInterval
	}
	log := opts.Logger
	if log == nil {
		log = slog.Default()
	}
	scale := opts.IntervalScale
	if scale <= 0 {
		scale = 1
	}
	sc := cfg.Schedule
	secs := func(v int, def time.Duration) time.Duration {
		if v <= 0 {
			return time.Duration(float64(def) * scale)
		}
		return time.Duration(float64(v) * scale * float64(time.Second))
	}
	clamp := func(d time.Duration, what string) time.Duration {
		if d < minI {
			log.Warn("interval below minimum, clamped", "what", what, "requested", d, "min", minI)
			return minI
		}
		return d
	}
	dnsEvery := clamp(secs(sc.DNSEvery, DefaultDNSEvery), "dns_every")
	pingEvery := clamp(secs(sc.PingEvery, DefaultPingEvery), "ping_every")
	trEvery := clamp(secs(sc.TracerouteEvery, DefaultTraceEvery), "traceroute_every")

	var jobs []Job
	for _, t := range cfg.Targets {
		if t.TargetID == "" || t.Hostname == "" {
			continue
		}
		tierI := secs(sc.Tiers[strconv.Itoa(t.Tier)], DefaultTierInterval)
		tierI = clamp(tierI, "tier "+strconv.Itoa(t.Tier))
		add := func(f Family, d time.Duration) {
			jobs = append(jobs, Job{Key: t.TargetID + "/" + string(f), Target: t, Family: f, Interval: d})
		}
		for _, c := range t.Checks {
			switch Family(c) {
			case FamilyHTTP:
				add(FamilyHTTP, tierI)
			case FamilyDNS:
				add(FamilyDNS, dnsEvery)
			case FamilyPing:
				if opts.CanPing {
					add(FamilyPing, pingEvery)
				} else {
					add(FamilyTCP, pingEvery) // icmp unavailable → tcp connects on the same cadence
				}
			case FamilyTCP:
				if !t.HasCheck("ping") || opts.CanPing { // avoid a duplicate tcp job when ping already fell back
					add(FamilyTCP, pingEvery)
				}
			}
		}
		if t.Traceroute && opts.CanTraceroute {
			add(FamilyTraceroute, trEvery)
		}
	}
	return jobs
}