package platform

import (
	"context"
	"encoding/csv"
	"fmt"
	"math"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"
)

func normalizeNvidiaBenchmarkOptionsForBenchmark(opts NvidiaBenchmarkOptions) NvidiaBenchmarkOptions {
	switch strings.TrimSpace(strings.ToLower(opts.Profile)) {
	case NvidiaBenchmarkProfileStability:
		opts.Profile = NvidiaBenchmarkProfileStability
	case NvidiaBenchmarkProfileOvernight:
		opts.Profile = NvidiaBenchmarkProfileOvernight
	default:
		opts.Profile = NvidiaBenchmarkProfileStandard
	}
	if opts.SizeMB < 0 {
		opts.SizeMB = 0
	}
	opts.ServerPowerSource = normalizeBenchmarkPowerSource(opts.ServerPowerSource)
	opts.GPUIndices = dedupeSortedIndices(opts.GPUIndices)
	opts.ExcludeGPUIndices = dedupeSortedIndices(opts.ExcludeGPUIndices)
	return opts
}

func resolveBenchmarkProfile(profile string) benchmarkProfileSpec {
	switch strings.TrimSpace(strings.ToLower(profile)) {
	case NvidiaBenchmarkProfileStability:
		return benchmarkProfileSpec{Name: NvidiaBenchmarkProfileStability, BaselineSec: 30, WarmupSec: 120, SteadySec: 3600, NCCLSec: 300, CooldownSec: 0}
	case NvidiaBenchmarkProfileOvernight:
		return benchmarkProfileSpec{Name: NvidiaBenchmarkProfileOvernight, BaselineSec: 60, WarmupSec: 180, SteadySec: 27000, NCCLSec: 600, CooldownSec: 0}
	default:
		return benchmarkProfileSpec{Name: NvidiaBenchmarkProfileStandard, BaselineSec: 15, WarmupSec: 45, SteadySec: 480, NCCLSec: 180, CooldownSec: 0}
	}
}

// benchmarkGPUInfoQuery describes a nvidia-smi --query-gpu field set to try.
// Fields are tried in order; the first successful query wins. Extended fields
// (attribute.multiprocessor_count, power.default_limit) are not supported on
// all driver versions, so we fall back to the base set if the full query fails.
// The minimal fallback omits clock fields entirely — clocks.max.* returns
// exit status 2 on some GPU generations (e.g. Blackwell); missing data is
// then recovered from nvidia-smi -q.
var benchmarkGPUInfoQueries = []struct {
	fields   string
	extended bool // whether this query includes optional extended fields
	minimal  bool // clock fields omitted; max clocks must be filled separately
}{
	{
		fields:   "index,uuid,name,pci.bus_id,vbios_version,power.limit,clocks.max.graphics,clocks.max.memory,clocks.base.graphics,attribute.multiprocessor_count,power.default_limit",
		extended: true,
	},
	{
		fields:   "index,uuid,name,pci.bus_id,vbios_version,power.limit,clocks.max.graphics,clocks.max.memory,clocks.base.graphics",
		extended: false,
	},
	{
		fields:  "index,uuid,name,pci.bus_id,vbios_version,power.limit",
		minimal: true,
	},
}

// enrichGPUInfoWithNvidiaSMIQ fills benchmark GPU metadata from nvidia-smi -q
// for fields that may be missing from --query-gpu on some driver versions.
func enrichGPUInfoWithNvidiaSMIQ(infoByIndex map[int]benchmarkGPUInfo, nvsmiQ []byte) {
	if len(infoByIndex) == 0 || len(nvsmiQ) == 0 {
		return
	}

	// Build bus_id → index map for matching verbose sections to GPU indices.
	busToBenchIdx := make(map[string]int, len(infoByIndex))
	for idx, info := range infoByIndex {
		if info.BusID != "" {
			// nvidia-smi -q uses "GPU 00000000:4E:00.0" (8-digit domain),
			// while --query-gpu returns the same format; normalise to lower.
			busToBenchIdx[strings.ToLower(strings.TrimSpace(info.BusID))] = idx
		}
	}

	// Split the verbose output into per-GPU sections on "^GPU " lines.
	gpuSectionRe := regexp.MustCompile(`(?m)^GPU\s+([\dA-Fa-f:\.]+)`)
	maxGfxRe := regexp.MustCompile(`(?i)Max Clocks[\s\S]*?Graphics\s*:\s*(\d+)\s*MHz`)
	maxMemRe := regexp.MustCompile(`(?i)Max Clocks[\s\S]*?Memory\s*:\s*(\d+)\s*MHz`)
	defaultPwrRe := regexp.MustCompile(`(?i)Default Power Limit\s*:\s*([0-9.]+)\s*W`)
	currentPwrRe := regexp.MustCompile(`(?i)Current Power Limit\s*:\s*([0-9.]+)\s*W`)
	minPwrRe := regexp.MustCompile(`(?i)Min Power Limit\s*:\s*([0-9.]+)\s*W`)
	maxPwrRe := regexp.MustCompile(`(?i)Max Power Limit\s*:\s*([0-9.]+)\s*W`)
	smCountRe := regexp.MustCompile(`(?i)Multiprocessor Count\s*:\s*(\d+)`)
	shutdownTempRe := regexp.MustCompile(`(?i)GPU Shutdown Temp\s*:\s*(\d+)\s*C`)
	slowdownTempRe := regexp.MustCompile(`(?i)GPU Slowdown Temp\s*:\s*(\d+)\s*C`)

	sectionStarts := gpuSectionRe.FindAllSubmatchIndex(nvsmiQ, -1)
	for i, loc := range sectionStarts {
		busID := strings.ToLower(string(nvsmiQ[loc[2]:loc[3]]))
		benchIdx, ok := busToBenchIdx[busID]
		if !ok {
			// Bus IDs from verbose output may have a different domain prefix;
			// try suffix match on the slot portion (XX:XX.X).
			for k, v := range busToBenchIdx {
				if strings.HasSuffix(k, busID) || strings.HasSuffix(busID, k) {
					benchIdx = v
					ok = true
					break
				}
			}
		}
		if !ok {
			continue
		}

		end := len(nvsmiQ)
		if i+1 < len(sectionStarts) {
			end = sectionStarts[i+1][0]
		}
		section := nvsmiQ[loc[0]:end]

		info := infoByIndex[benchIdx]

		if info.MaxGraphicsClockMHz == 0 {
			if m := maxGfxRe.FindSubmatch(section); m != nil {
				if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil {
					info.MaxGraphicsClockMHz = v
				}
			}
		}
		if info.MaxMemoryClockMHz == 0 {
			if m := maxMemRe.FindSubmatch(section); m != nil {
				if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil {
					info.MaxMemoryClockMHz = v
				}
			}
		}
		if info.DefaultPowerLimitW == 0 {
			if m := defaultPwrRe.FindSubmatch(section); m != nil {
				if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
					info.DefaultPowerLimitW = v
				}
			}
		}
		if info.PowerLimitW == 0 {
			if m := currentPwrRe.FindSubmatch(section); m != nil {
				if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
					info.PowerLimitW = v
				}
			}
		}
		if info.MinPowerLimitW == 0 {
			if m := minPwrRe.FindSubmatch(section); m != nil {
				if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
					info.MinPowerLimitW = v
				}
			}
		}
		if info.MaxPowerLimitW == 0 {
			if m := maxPwrRe.FindSubmatch(section); m != nil {
				if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
					info.MaxPowerLimitW = v
				}
			}
		}
		if info.MultiprocessorCount == 0 {
			if m := smCountRe.FindSubmatch(section); m != nil {
				if v, err := strconv.Atoi(string(m[1])); err == nil && v > 0 {
					info.MultiprocessorCount = v
				}
			}
		}
		if info.ShutdownTempC == 0 {
			if m := shutdownTempRe.FindSubmatch(section); m != nil {
				if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
					info.ShutdownTempC = v
				}
			}
		}
		if info.SlowdownTempC == 0 {
			if m := slowdownTempRe.FindSubmatch(section); m != nil {
				if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
					info.SlowdownTempC = v
				}
			}
		}
		infoByIndex[benchIdx] = info
	}
}

func queryBenchmarkGPUInfo(gpuIndices []int) (map[int]benchmarkGPUInfo, error) {
	var lastErr error
	for _, q := range benchmarkGPUInfoQueries {
		args := []string{
			"--query-gpu=" + q.fields,
			"--format=csv,noheader,nounits",
		}
		if len(gpuIndices) > 0 {
			args = append([]string{"--id=" + joinIndexList(gpuIndices)}, args...)
		}
		out, err := satExecCommand("nvidia-smi", args...).Output()
		if err != nil {
			lastErr = fmt.Errorf("nvidia-smi gpu info (%s): %w", q.fields[:min(len(q.fields), 40)], err)
			continue
		}

		r := csv.NewReader(strings.NewReader(string(out)))
		r.TrimLeadingSpace = true
		r.FieldsPerRecord = -1
		rows, err := r.ReadAll()
		if err != nil {
			lastErr = fmt.Errorf("parse nvidia-smi gpu info: %w", err)
			continue
		}

		minFields := 6
		if !q.minimal {
			minFields = 9
		}
		infoByIndex := make(map[int]benchmarkGPUInfo, len(rows))
		for _, row := range rows {
			if len(row) < minFields {
				continue
			}
			idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
			if err != nil {
				continue
			}
			info := benchmarkGPUInfo{
				Index:       idx,
				UUID:        strings.TrimSpace(row[1]),
				Name:        strings.TrimSpace(row[2]),
				BusID:       strings.TrimSpace(row[3]),
				VBIOS:       strings.TrimSpace(row[4]),
				PowerLimitW: parseBenchmarkFloat(row[5]),
			}
			if !q.minimal {
				info.MaxGraphicsClockMHz = parseBenchmarkFloat(row[6])
				info.MaxMemoryClockMHz = parseBenchmarkFloat(row[7])
				if len(row) >= 9 {
					info.BaseGraphicsClockMHz = parseBenchmarkFloat(row[8])
				}
				if q.extended {
					if len(row) >= 10 {
						info.MultiprocessorCount = int(parseBenchmarkFloat(row[9]))
					}
					if len(row) >= 11 {
						info.DefaultPowerLimitW = parseBenchmarkFloat(row[10])
					}
				}
			}
			infoByIndex[idx] = info
		}
		return infoByIndex, nil
	}
	return nil, lastErr
}

func applyBenchmarkNormalization(ctx context.Context, verboseLog string, gpuIndices []int, infoByIndex map[int]benchmarkGPUInfo, result *NvidiaBenchmarkResult) []benchmarkRestoreAction {
	if os.Geteuid() != 0 {
		result.Normalization.Status = "partial"
		result.Normalization.Notes = append(result.Normalization.Notes, "benchmark normalization skipped: root privileges are required for persistence mode and clock locks")
		for _, idx := range gpuIndices {
			result.Normalization.GPUs = append(result.Normalization.GPUs, BenchmarkNormalizationGPU{
				Index: idx,
				Notes: []string{"normalization skipped: root privileges are required"},
			})
		}
		return nil
	}

	var restore []benchmarkRestoreAction
	for _, idx := range gpuIndices {
		rec := BenchmarkNormalizationGPU{Index: idx}
		if _, err := runSATCommandCtx(ctx, verboseLog, fmt.Sprintf("normalize-gpu-%d-pm", idx), []string{"nvidia-smi", "-i", strconv.Itoa(idx), "-pm", "1"}, nil, nil); err != nil {
			rec.PersistenceMode = "failed"
			rec.Notes = append(rec.Notes, "failed to enable persistence mode")
			result.Normalization.Status = "partial"
		} else {
			rec.PersistenceMode = "applied"
		}

		if info, ok := infoByIndex[idx]; ok && info.MaxGraphicsClockMHz > 0 {
			target := int(math.Round(info.MaxGraphicsClockMHz))
			if out, err := runSATCommandCtx(ctx, verboseLog, fmt.Sprintf("normalize-gpu-%d-lgc", idx), []string{"nvidia-smi", "-i", strconv.Itoa(idx), "-lgc", strconv.Itoa(target)}, nil, nil); err != nil {
				rec.GPUClockLockStatus = "failed"
				rec.Notes = append(rec.Notes, "graphics clock lock failed: "+strings.TrimSpace(string(out)))
				result.Normalization.Status = "partial"
			} else {
				rec.GPUClockLockStatus = "applied"
				rec.GPUClockLockMHz = float64(target)
				idxCopy := idx
				restore = append(restore, benchmarkRestoreAction{name: fmt.Sprintf("gpu-%d-rgc", idxCopy), fn: func() {
					_, _ = runSATCommandCtx(context.Background(), verboseLog, fmt.Sprintf("restore-gpu-%d-rgc", idxCopy), []string{"nvidia-smi", "-i", strconv.Itoa(idxCopy), "-rgc"}, nil, nil)
				}})
			}
		} else {
			rec.GPUClockLockStatus = "skipped"
			rec.Notes = append(rec.Notes, "graphics clock lock skipped: gpu inventory unavailable or MaxGraphicsClockMHz=0")
			result.Normalization.Status = "partial"
		}

		if info, ok := infoByIndex[idx]; ok && info.MaxMemoryClockMHz > 0 {
			target := int(math.Round(info.MaxMemoryClockMHz))
			out, err := runSATCommandCtx(ctx, verboseLog, fmt.Sprintf("normalize-gpu-%d-lmc", idx), []string{"nvidia-smi", "-i", strconv.Itoa(idx), "-lmc", strconv.Itoa(target)}, nil, nil)
			switch {
			case err == nil:
				rec.MemoryClockLockStatus = "applied"
				rec.MemoryClockLockMHz = float64(target)
				idxCopy := idx
				restore = append(restore, benchmarkRestoreAction{name: fmt.Sprintf("gpu-%d-rmc", idxCopy), fn: func() {
					_, _ = runSATCommandCtx(context.Background(), verboseLog, fmt.Sprintf("restore-gpu-%d-rmc", idxCopy), []string{"nvidia-smi", "-i", strconv.Itoa(idxCopy), "-rmc"}, nil, nil)
				}})
			case strings.Contains(strings.ToLower(string(out)), "deferred") || strings.Contains(strings.ToLower(string(out)), "not supported"):
				rec.MemoryClockLockStatus = "unsupported"
				rec.Notes = append(rec.Notes, "memory clock lock unsupported on this GPU/driver path")
				result.Normalization.Status = "partial"
			default:
				rec.MemoryClockLockStatus = "failed"
				rec.Notes = append(rec.Notes, "memory clock lock failed: "+strings.TrimSpace(string(out)))
				result.Normalization.Status = "partial"
			}
		}

		result.Normalization.GPUs = append(result.Normalization.GPUs, rec)
	}
	return restore
}

func collectBenchmarkSamples(ctx context.Context, durationSec int, gpuIndices []int) ([]GPUMetricRow, error) {
	if durationSec <= 0 {
		return nil, nil
	}
	deadline := time.Now().Add(time.Duration(durationSec) * time.Second)
	var rows []GPUMetricRow
	start := time.Now()
	for {
		if ctx.Err() != nil {
			return rows, ctx.Err()
		}
		samples, err := sampleBenchmarkTelemetry(gpuIndices)
		if err == nil {
			elapsed := time.Since(start).Seconds()
			for i := range samples {
				samples[i].ElapsedSec = elapsed
			}
			rows = append(rows, samples...)
		}
		if time.Now().After(deadline) {
			break
		}
		select {
		case <-ctx.Done():
			return rows, ctx.Err()
		case <-time.After(time.Second):
		}
	}
	return rows, nil
}

func runBenchmarkCommandWithMetrics(ctx context.Context, verboseLog, name string, cmd []string, env []string, gpuIndices []int, logFunc func(string)) ([]byte, []GPUMetricRow, error) {
	stopCh := make(chan struct{})
	doneCh := make(chan struct{})
	var metricRows []GPUMetricRow
	start := time.Now()

	go func() {
		defer close(doneCh)
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-stopCh:
				return
			case <-ticker.C:
				samples, err := sampleBenchmarkTelemetry(gpuIndices)
				if err != nil {
					continue
				}
				elapsed := time.Since(start).Seconds()
				for i := range samples {
					samples[i].ElapsedSec = elapsed
				}
				metricRows = append(metricRows, samples...)
			}
		}
	}()

	out, err := runSATCommandCtx(ctx, verboseLog, name, cmd, env, logFunc)
	close(stopCh)
	<-doneCh

	return out, metricRows, err
}