package scheduler

// This file demonstrates how to add a custom scheduler.
// To create your own:
//   1. Copy this file and rename it (e.g., weighted_rtt.go)
//   2. Implement the three PathScheduler methods
//   3. Update the init() to register with your scheduler name
//
// The scheduler will automatically appear in --list-schedulers.

import (
	"sync"

	quic "github.com/AeonDave/mp-quic-go"
)

func init() {
	Register("weighted", func() quic.PathScheduler {
		return NewWeightedScheduler()
	})
}

// WeightedScheduler is an example custom scheduler that weighs paths
// by a combination of RTT and available congestion window.
type WeightedScheduler struct {
	mu     sync.Mutex
	quotas map[quic.PathID]uint64
}

// NewWeightedScheduler creates a new weighted scheduler.
func NewWeightedScheduler() *WeightedScheduler {
	return &WeightedScheduler{
		quotas: make(map[quic.PathID]uint64),
	}
}

// SelectPath picks the path with the best weighted score.
// Score = (CongestionWindow - BytesInFlight) / (1 + SmoothedRTT_ms)
// Higher score = more capacity available per unit latency.
func (s *WeightedScheduler) SelectPath(paths []quic.SchedulerPathInfo, hasRetransmission bool) *quic.SchedulerPathInfo {
	s.mu.Lock()
	defer s.mu.Unlock()

	if len(paths) == 0 {
		return nil
	}
	if len(paths) == 1 {
		if !hasRetransmission && !paths[0].SendingAllowed {
			return nil
		}
		return &paths[0]
	}

	var best *quic.SchedulerPathInfo
	var bestScore float64 = -1

	for i := range paths {
		p := &paths[i]
		if !hasRetransmission && !p.SendingAllowed {
			continue
		}
		if p.PotentiallyFailed {
			continue
		}

		// Available window
		available := float64(0)
		if p.CongestionWindow > p.BytesInFlight {
			available = float64(p.CongestionWindow - p.BytesInFlight)
		}

		// RTT factor (ms, minimum 1 to avoid division by zero)
		rttMs := float64(1)
		if p.SmoothedRTT.Milliseconds() > 0 {
			rttMs = float64(p.SmoothedRTT.Milliseconds())
		}

		score := available / rttMs
		if score > bestScore {
			bestScore = score
			best = p
		}
	}

	return best
}

// UpdateQuota tracks per-path send counts.
func (s *WeightedScheduler) UpdateQuota(pathID quic.PathID, packetSize quic.ByteCount) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.quotas[pathID]++
}

// Reset clears all state.
func (s *WeightedScheduler) Reset() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.quotas = make(map[quic.PathID]uint64)
}