package platform

import (
	"context"
	"errors"
	"fmt"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

func runBenchmarkPowerCalibration(
	ctx context.Context,
	verboseLog, runDir string,
	gpuIndices []int,
	infoByIndex map[int]benchmarkGPUInfo,
	logFunc func(string),
	seedLimits map[int]int,
	durationSec int,
) (map[int]benchmarkPowerCalibrationResult, []benchmarkRestoreAction, []GPUMetricRow, benchmarkPowerCalibrationRunSummary) {
	calibDurationSec := durationSec
	var runSummary benchmarkPowerCalibrationRunSummary
	if calibDurationSec <= 0 {
		calibDurationSec = 120
	}
	// calibSearchTolerance is the binary-search convergence threshold in watts.
	// When hi-lo ≤ this, the highest verified-stable limit (lo) is used.
	const calibSearchTolerance = 10
	// dcgmResourceBusyMaxDelaySec caps the exponential back-off when DCGM
	// returns DCGM_ST_IN_USE (exit 222). The sequence is 1 s, 2 s, 4 s, …
	// doubling each retry until it would exceed the cap, at which point the
	// next busy response fails the calibration immediately.
	const dcgmResourceBusyMaxDelaySec = 300
	engine := benchmarkPowerEngine()
	engineLabel := benchmarkPowerEngineLabel(engine)

	if engine == BenchmarkPowerEngineTargetedPower {
		if _, err := exec.LookPath("dcgmi"); err != nil {
			logFunc("power calibration: dcgmi not found, skipping (will use default power limit)")
			return map[int]benchmarkPowerCalibrationResult{}, nil, nil, runSummary
		}
	} else {
		if _, _, err := resolveBenchmarkPowerLoadCommand(calibDurationSec, gpuIndices); err != nil {
			logFunc("power calibration: dcgmproftester not found, skipping (will use default power limit)")
			return map[int]benchmarkPowerCalibrationResult{}, nil, nil, runSummary
		}
	}
	if killed := KillTestWorkers(); len(killed) > 0 {
		for _, p := range killed {
			logFunc(fmt.Sprintf("power calibration pre-flight: killed stale worker pid=%d name=%s", p.PID, p.Name))
		}
	}

	canDerate := os.Geteuid() == 0
	if !canDerate {
		logFunc("power calibration: root privileges unavailable, adaptive power-limit derating disabled")
	}

	type calibrationAttemptResult struct {
		out  []byte
		rows []GPUMetricRow
		err  error
	}

	// gpuCalibState holds per-GPU binary search state during parallel calibration.
	type gpuCalibState struct {
		idx            int
		info           benchmarkGPUInfo
		originalLimitW int
		appliedLimitW  int
		minLimitW      int
		lo             int // highest verified-stable limit
		hi             int // lowest verified-unstable limit (exclusive sentinel above start)
		loVerified     bool
		calib          benchmarkPowerCalibrationResult
		converged      bool
	}

	results := make(map[int]benchmarkPowerCalibrationResult, len(gpuIndices))
	var restore []benchmarkRestoreAction
	var allCalibRows []GPUMetricRow // accumulated telemetry across all attempts
	var calibCursor float64

	// Initialise per-GPU state.
	states := make([]*gpuCalibState, 0, len(gpuIndices))
	for _, idx := range gpuIndices {
		info := infoByIndex[idx]
		originalLimitW := int(math.Round(info.PowerLimitW))
		if originalLimitW <= 0 {
			originalLimitW = int(math.Round(info.DefaultPowerLimitW))
		}
		defaultLimitW := int(math.Round(info.DefaultPowerLimitW))
		if defaultLimitW <= 0 {
			defaultLimitW = originalLimitW
		}
		appliedLimitW := initialBenchmarkCalibrationLimitW(info)
		if appliedLimitW <= 0 {
			appliedLimitW = defaultLimitW
		}
		minLimitW := int(math.Round(info.MinPowerLimitW))
		if minLimitW <= 0 {
			minLimitW = appliedLimitW
		}
		maxLimitW := int(math.Round(info.MaxPowerLimitW))
		if maxLimitW > 0 && appliedLimitW > maxLimitW {
			appliedLimitW = maxLimitW
		}
		s := &gpuCalibState{
			idx:            idx,
			info:           info,
			originalLimitW: originalLimitW,
			appliedLimitW:  appliedLimitW,
			minLimitW:      minLimitW,
			lo:             minLimitW,
			hi:             appliedLimitW + 1, // not yet tested, not yet confirmed unstable
			calib:          benchmarkPowerCalibrationResult{AppliedPowerLimitW: float64(appliedLimitW)},
		}
		if minLimitW > 0 && appliedLimitW > 0 && minLimitW >= appliedLimitW {
			s.appliedLimitW = minLimitW
			s.hi = minLimitW + 1
		}
		if info.MinPowerLimitW <= 0 {
			s.calib.Notes = append(s.calib.Notes, "minimum power limit was not reported by nvidia-smi; calibration can only validate the current/default power limit")
		}
		if seedLimits != nil {
			if seedW, ok := seedLimits[idx]; ok && seedW > 0 {
				// A previously validated limit is only a starting point. Re-run
				// targeted_power under the current multi-GPU thermal load and derate
				// again if this step shows new throttling.
				if seedW < s.minLimitW {
					seedW = s.minLimitW
				}
				if maxLimitW > 0 && seedW > maxLimitW {
					seedW = maxLimitW
				}
				if canDerate {
					_ = setBenchmarkPowerLimit(ctx, verboseLog, idx, seedW)
				}
				s.appliedLimitW = seedW
				s.hi = seedW + 1
				s.calib.AppliedPowerLimitW = float64(seedW)
				s.calib.Derated = seedW < s.originalLimitW
				s.calib.Notes = append(s.calib.Notes,
					fmt.Sprintf("seed limit: %d W (revalidating under current thermal load)", seedW))
			}
		}
		states = append(states, s)
		if canDerate && originalLimitW > 0 {
			idxCopy := idx
			orig := originalLimitW
			restore = append(restore, benchmarkRestoreAction{
				name: fmt.Sprintf("gpu-%d-restore-power-limit", idxCopy),
				fn: func() {
					_ = setBenchmarkPowerLimit(context.Background(), verboseLog, idxCopy, orig)
				},
			})
		}
	}

	// Shared DCGM resource-busy back-off state (single diagnostic session).
	busyRetries := 0
	busyDelaySec := 1
	sharedAttempt := 0

	type sharedAttemptResult struct {
		out  []byte
		rows []GPUMetricRow
		err  error
	}

calibDone:
	for {
		// Collect non-converged GPUs.
		var active []*gpuCalibState
		for _, s := range states {
			if !s.converged {
				active = append(active, s)
			}
		}
		if len(active) == 0 || ctx.Err() != nil {
			break
		}

		sharedAttempt++
		for _, s := range active {
			s.calib.Attempts++
			logFunc(fmt.Sprintf("power calibration: GPU %d %s attempt %d at %d W for %ds", s.idx, engineLabel, s.calib.Attempts, s.appliedLimitW, calibDurationSec))
		}

		// Snapshot throttle counters for all active GPUs before the run.
		beforeThrottle := make(map[int]BenchmarkThrottleCounters, len(active))
		for _, s := range active {
			beforeThrottle[s.idx], _ = queryThrottleCounters(s.idx)
		}

		// Run the selected power-fit load for ALL gpuIndices simultaneously so every card
		// is under load during calibration — this reflects real server thermals.
		logName := fmt.Sprintf("power-calibration-attempt-%d.log", sharedAttempt)
		cmd, env, err := resolveBenchmarkPowerLoadCommand(calibDurationSec, gpuIndices)
		if err != nil {
			for _, s := range active {
				s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("failed to resolve %s command: %v", engineLabel, err))
				s.converged = true
			}
			logFunc(fmt.Sprintf("power calibration: failed to resolve %s command: %v", engineLabel, err))
			break calibDone
		}
		attemptCtx, cancelAttempt := context.WithCancel(ctx)
		doneCh := make(chan sharedAttemptResult, 1)
		sdrStopCh := make(chan struct{})
		sdrDoneCh := startIPMISDRSampler(sdrStopCh, benchmarkPowerAutotuneSampleInterval)
		fanStopCh := make(chan struct{})
		fanDoneCh := startBenchmarkFanSampler(fanStopCh, benchmarkPowerAutotuneSampleInterval)
		go func() {
			out, rows, err := runBenchmarkCommandWithMetrics(attemptCtx, verboseLog, logName, cmd, env, gpuIndices, logFunc)
			doneCh <- sharedAttemptResult{out: out, rows: rows, err: err}
		}()

		ticker := time.NewTicker(time.Second)
		throttleReasons := make(map[int]string, len(active))
		var ar sharedAttemptResult

	attemptLoop:
		for {
			select {
			case ar = <-doneCh:
				break attemptLoop
			case <-ticker.C:
				// Poll throttle counters for each active GPU independently.
				for _, s := range active {
					if throttleReasons[s.idx] != "" {
						continue // already detected for this GPU
					}
					after, err := queryThrottleCounters(s.idx)
					if err != nil {
						continue
					}
					// Record throttle but do NOT cancel — let the load command finish so
					// runtime resources release cleanly before the next attempt.
					if reason := benchmarkCalibrationThrottleReason(beforeThrottle[s.idx], after); reason != "" {
						throttleReasons[s.idx] = reason
						logFunc(fmt.Sprintf("power calibration: GPU %d detected %s throttle at %d W, waiting for run to finish", s.idx, reason, s.appliedLimitW))
					}
				}
			case <-ctx.Done():
				cancelAttempt()
				ar = <-doneCh
				break attemptLoop
			}
		}
		ticker.Stop()
		cancelAttempt()
		close(sdrStopCh)
		close(fanStopCh)
		attemptSDRSummary := summarizeSDRPowerSeries(<-sdrDoneCh)
		attemptFanSummary := <-fanDoneCh
		_ = os.WriteFile(filepath.Join(runDir, logName), ar.out, 0644)
		// Accumulate telemetry rows with attempt stage label.
		appendBenchmarkMetrics(&allCalibRows, ar.rows, fmt.Sprintf("attempt-%d", sharedAttempt), &calibCursor, float64(calibDurationSec))

		// Resource busy: retry with exponential back-off (shared — one DCGM session).
		if ar.err != nil && isDCGMResourceBusy(ar.err) {
			if busyDelaySec > dcgmResourceBusyMaxDelaySec {
				for _, s := range active {
					s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("DCGM resource busy after %d retries, giving up", busyRetries))
					s.converged = true
				}
				logFunc(fmt.Sprintf("power calibration: DCGM resource persistently busy after %d retries, stopping", busyRetries))
				break calibDone
			}
			busyRetries++
			// Undo attempt counter: busy retries don't count as real attempts.
			for _, s := range active {
				s.calib.Attempts--
			}
			logFunc(fmt.Sprintf("power calibration: DCGM resource busy (attempt %d), retrying in %ds", sharedAttempt, busyDelaySec))
			select {
			case <-ctx.Done():
				break calibDone
			case <-time.After(time.Duration(busyDelaySec) * time.Second):
			}
			next := busyDelaySec * 2
			if next > dcgmResourceBusyMaxDelaySec {
				next = dcgmResourceBusyMaxDelaySec + 1
			}
			busyDelaySec = next
			sharedAttempt-- // retry same logical attempt number
			continue
		}
		busyRetries = 0
		busyDelaySec = 1

		// Per-GPU analysis and binary search update.
		attemptStable := ar.err == nil
		for _, s := range active {
			perGPU := filterRowsByGPU(ar.rows, s.idx)
			summary := summarizeBenchmarkTelemetry(perGPU)
			throttle := throttleReasons[s.idx]
			if throttle != "" || summary.P95PowerW <= 0 {
				attemptStable = false
			}

			// Cooling warning: thermal throttle with fans not at maximum.
			if strings.Contains(throttle, "thermal") && s.calib.CoolingWarning == "" {
				clocks := make([]float64, 0, len(perGPU))
				var fanDutyValues []float64
				fanDutyAvail := false
				for _, r := range perGPU {
					if r.ClockMHz > 0 {
						clocks = append(clocks, r.ClockMHz)
					}
					if r.FanDutyCycleAvailable {
						fanDutyAvail = true
						fanDutyValues = append(fanDutyValues, r.FanDutyCyclePct)
					}
				}
				dropPct := benchmarkClockDrift(clocks)
				p95FanDuty := benchmarkPercentile(fanDutyValues, 95)
				if dropPct >= 20 && fanDutyAvail && p95FanDuty < 98 {
					s.calib.CoolingWarning = fmt.Sprintf(
						"thermal throttle (%s) caused a %.0f%% clock drop while fans were at %.0f%% duty cycle — server cooling may not be configured for full GPU load",
						throttle, dropPct, p95FanDuty,
					)
					logFunc(fmt.Sprintf("power calibration: GPU %d cooling warning: %s", s.idx, s.calib.CoolingWarning))
				}
			}

			if throttle == "" && ar.err == nil && summary.P95PowerW > 0 {
				// Stable at current limit — update lo and binary-search upward.
				s.calib.Summary = summary
				s.calib.Completed = true
				s.calib.AppliedPowerLimitW = float64(s.appliedLimitW)
				logFunc(fmt.Sprintf("power calibration: GPU %d stable at %d W, p95=%.0f W p95_temp=%.1f C (%d samples)", s.idx, s.appliedLimitW, summary.P95PowerW, summary.P95TempC, summary.Samples))
				s.lo = s.appliedLimitW
				s.loVerified = true
				if canDerate && s.hi-s.lo > calibSearchTolerance {
					next := roundTo5W((s.lo + s.hi) / 2)
					if next > s.lo && next < s.hi {
						if err := setBenchmarkPowerLimit(ctx, verboseLog, s.idx, next); err == nil {
							s.appliedLimitW = next
							s.calib.AppliedPowerLimitW = float64(next)
							s.calib.Completed = false // keep searching
							s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("binary search: stable at %d W, trying %d W (lo=%d hi=%d)", s.lo, next, s.lo, s.hi))
							logFunc(fmt.Sprintf("power calibration: GPU %d binary search up: stable at %d W, trying %d W", s.idx, s.lo, next))
							continue // next GPU in active list
						}
					}
				}
				s.calib.MetricRows = filterRowsByGPU(ar.rows, s.idx)
				s.converged = true
				continue
			}

			// Failed or throttled — log and binary-search downward.
			switch {
			case throttle != "":
				s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("targeted_power attempt %d: %s throttle at %d W", s.calib.Attempts, throttle, s.appliedLimitW))
				logFunc(fmt.Sprintf("power calibration: GPU %d throttled (%s) at %d W, reducing power limit", s.idx, throttle, s.appliedLimitW))
			case ar.err != nil:
				s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("targeted_power attempt %d failed at %d W: %v", s.calib.Attempts, s.appliedLimitW, ar.err))
				logFunc(fmt.Sprintf("power calibration: GPU %d %s failed at %d W: %v", s.idx, engineLabel, s.appliedLimitW, ar.err))
			default:
				s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("%s attempt %d at %d W: no valid power telemetry", engineLabel, s.calib.Attempts, s.appliedLimitW))
				logFunc(fmt.Sprintf("power calibration: GPU %d attempt %d at %d W: no valid telemetry", s.idx, s.calib.Attempts, s.appliedLimitW))
			}

			if !canDerate || s.appliedLimitW <= 0 {
				s.converged = true
				continue
			}
			s.hi = s.appliedLimitW

			if s.hi-s.lo <= calibSearchTolerance {
				if !s.loVerified && s.minLimitW > 0 && s.appliedLimitW != s.minLimitW {
					if err := setBenchmarkPowerLimit(ctx, verboseLog, s.idx, s.minLimitW); err != nil {
						s.calib.Notes = append(s.calib.Notes, "failed to set power limit: "+err.Error())
						logFunc(fmt.Sprintf("power calibration: GPU %d failed to set minimum power limit %d W: %v", s.idx, s.minLimitW, err))
						s.converged = true
						continue
					}
					s.appliedLimitW = s.minLimitW
					s.calib.AppliedPowerLimitW = float64(s.minLimitW)
					s.calib.Derated = s.minLimitW < s.originalLimitW
					s.info.PowerLimitW = float64(s.minLimitW)
					infoByIndex[s.idx] = s.info
					s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("binary search: validating minimum settable limit %d W before concluding failure", s.minLimitW))
					logFunc(fmt.Sprintf("power calibration: GPU %d binary search: validating minimum settable limit %d W", s.idx, s.minLimitW))
					continue
				}
				if s.loVerified {
					s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("binary search converged: using %d W (lo=%d hi=%d)", s.lo, s.lo, s.hi))
					if err := setBenchmarkPowerLimit(ctx, verboseLog, s.idx, s.lo); err == nil {
						s.appliedLimitW = s.lo
						s.calib.AppliedPowerLimitW = float64(s.lo)
						s.calib.Derated = s.lo < s.originalLimitW
						// Summary was captured when we last verified stability at s.lo,
						// so the result is valid — mark as completed even though we
						// converged from the failure path (tried higher, failed, fell back).
						s.calib.Completed = true
					}
				} else {
					s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("could not find a stable %s limit down to the minimum settable power limit %d W", engineLabel, s.minLimitW))
					logFunc(fmt.Sprintf("power calibration: GPU %d no stable limit found down to minimum settable power limit %d W", s.idx, s.minLimitW))
				}
				s.calib.MetricRows = filterRowsByGPU(ar.rows, s.idx)
				s.converged = true
				continue
			}

			next := roundTo5W((s.lo + s.hi) / 2)
			if next <= s.lo {
				next = s.lo + calibSearchTolerance
			}
			if next >= s.hi {
				next = (s.lo + s.hi) / 2
			}
			if next < s.minLimitW {
				next = s.minLimitW
			}
			if err := setBenchmarkPowerLimit(ctx, verboseLog, s.idx, next); err != nil {
				s.calib.Notes = append(s.calib.Notes, "failed to set power limit: "+err.Error())
				logFunc(fmt.Sprintf("power calibration: GPU %d failed to set power limit %d W: %v", s.idx, next, err))
				s.converged = true
				continue
			}
			s.appliedLimitW = next
			s.calib.AppliedPowerLimitW = float64(next)
			s.calib.Derated = next < s.originalLimitW
			s.info.PowerLimitW = float64(next)
			infoByIndex[s.idx] = s.info
			s.calib.Notes = append(s.calib.Notes, fmt.Sprintf("binary search: trying %d W (lo=%d hi=%d)", next, s.lo, s.hi))
			logFunc(fmt.Sprintf("power calibration: GPU %d binary search: trying %d W (lo=%d hi=%d)", s.idx, next, s.lo, s.hi))
		}
		if attemptStable {
			if attemptSDRSummary.Samples > 0 {
				runSummary.LoadedSDR = attemptSDRSummary
			}
			if attemptFanSummary.FanSamples > 0 {
				runSummary.AvgFanRPM = attemptFanSummary.AvgFanRPM
				runSummary.AvgFanDutyCyclePct = attemptFanSummary.AvgFanDutyCyclePct
				runSummary.FanSamples = attemptFanSummary.FanSamples
			}
		}
	}

	for _, s := range states {
		if s.calib.Completed || s.calib.Attempts > 0 || len(s.calib.Notes) > 0 {
			results[s.idx] = s.calib
		}
	}
	writeBenchmarkMetricsFiles(runDir, allCalibRows)
	return results, restore, allCalibRows, runSummary
}

// isDCGMResourceBusy returns true when dcgmi exits with DCGM_ST_IN_USE (222),
// meaning nv-hostengine still holds the diagnostic slot from a prior run.
func isDCGMResourceBusy(err error) bool {
	var exitErr *exec.ExitError
	return errors.As(err, &exitErr) && exitErr.ExitCode() == 222
}

// roundTo5W rounds w to the nearest 5 W boundary.
func roundTo5W(w int) int {
	return ((w + 2) / 5) * 5
}

func initialBenchmarkCalibrationLimitW(info benchmarkGPUInfo) int {
	defaultLimitW := int(math.Round(info.DefaultPowerLimitW))
	currentLimitW := int(math.Round(info.PowerLimitW))
	maxLimitW := int(math.Round(info.MaxPowerLimitW))

	startW := defaultLimitW
	if startW <= 0 {
		startW = currentLimitW
	}
	if startW <= 0 {
		startW = maxLimitW
	}
	if maxLimitW > 0 && startW > maxLimitW {
		startW = maxLimitW
	}
	return startW
}

// meanFanRPM returns the average RPM across a set of fan readings.
func meanFanRPM(fans []FanReading) float64 {
	if len(fans) == 0 {
		return 0
	}
	var sum float64
	for _, f := range fans {
		sum += f.RPM
	}
	return sum / float64(len(fans))
}

func startBenchmarkFanSampler(stopCh <-chan struct{}, intervalSec int) <-chan benchmarkPowerCalibrationRunSummary {
	if intervalSec <= 0 {
		intervalSec = benchmarkPowerAutotuneSampleInterval
	}
	ch := make(chan benchmarkPowerCalibrationRunSummary, 1)
	go func() {
		defer close(ch)
		var rpmSamples []float64
		var dutySamples []float64
		record := func() {
			fans, err := sampleFanSpeeds()
			if err != nil || len(fans) == 0 {
				return
			}
			if rpm := meanFanRPM(fans); rpm > 0 {
				rpmSamples = append(rpmSamples, rpm)
			}
			if duty, ok, _ := sampleFanDutyCyclePctFromFans(fans); ok && duty > 0 {
				dutySamples = append(dutySamples, duty)
			}
		}
		record()
		ticker := time.NewTicker(time.Duration(intervalSec) * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-stopCh:
				ch <- benchmarkPowerCalibrationRunSummary{
					AvgFanRPM:          benchmarkMean(rpmSamples),
					AvgFanDutyCyclePct: benchmarkMean(dutySamples),
					FanSamples:         len(rpmSamples),
				}
				return
			case <-ticker.C:
				record()
			}
		}
	}()
	return ch
}

func powerBenchDurationSec(profile string) int {
	switch strings.TrimSpace(strings.ToLower(profile)) {
	case NvidiaBenchmarkProfileStability:
		return 300
	case NvidiaBenchmarkProfileOvernight:
		return 600
	default:
		return 120
	}
}

func cloneBenchmarkGPUInfoMap(src map[int]benchmarkGPUInfo) map[int]benchmarkGPUInfo {
	out := make(map[int]benchmarkGPUInfo, len(src))
	for k, v := range src {
		out[k] = v
	}
	return out
}