package gpu

import (
	"context"
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/1Panel-dev/1Panel/agent/utils/cmd"
)

const (
	amdSMICommand     = "amd-smi"
	amdSMIDefaultPath = "/opt/rocm/bin/amd-smi"
)

type amdSMI struct {
	command string
}

func findAMDSMI() (string, bool) {
	for _, command := range []string{amdSMICommand, amdSMIDefaultPath} {
		if cmd.Which(command) {
			return command, true
		}
	}
	return "", false
}

func (a amdSMI) LoadInfo(ctx context.Context) (*Info, error) {
	var (
		staticData   string
		metricData   string
		extendedData string
		extendedErr  error
		processData  string
		staticErr    error
		metricErr    error
		processErr   error
		wg           sync.WaitGroup
	)

	wg.Add(4)
	go func() {
		defer wg.Done()
		extendedData, extendedErr = runAMDSMI(ctx, a.command, "metric", "--clock", "--ecc", "--json")
	}()
	go func() {
		defer wg.Done()
		staticData, staticErr = runAMDSMI(ctx, a.command, "static", "--asic", "--bus", "--driver", "--limit", "--json")
	}()
	go func() {
		defer wg.Done()
		metricData, metricErr = runAMDSMI(ctx, a.command, "metric", "--usage", "--power", "--temperature", "--mem-usage", "--fan", "--perf-level", "--json")
	}()
	go func() {
		defer wg.Done()
		processData, processErr = runAMDSMI(ctx, a.command, "process", "--general", "--json")
	}()
	wg.Wait()

	if staticErr != nil {
		return nil, fmt.Errorf("calling %s static failed: %w", a.command, staticErr)
	}
	info, err := parseAMDStatic(staticData)
	if err != nil {
		return nil, fmt.Errorf("parsing %s static output failed: %w", a.command, err)
	}

	if metricErr == nil {
		metricErr = applyAMDMetrics(info, metricData)
	}
	if metricErr != nil {
		info.Warnings = append(info.Warnings, fmt.Sprintf("%s metrics: %v", a.command, metricErr))
	}
	if extendedErr == nil {
		extendedErr = applyAMDMetrics(info, extendedData)
	}
	if extendedErr != nil {
		info.Warnings = append(info.Warnings, fmt.Sprintf("%s extended metrics: %v", a.command, extendedErr))
	}
	if processErr == nil {
		processErr = applyAMDProcesses(info, processData)
	}
	if processErr != nil {
		info.Warnings = append(info.Warnings, fmt.Sprintf("%s processes: %v", a.command, processErr))
	}

	return info, nil
}

func runAMDSMI(ctx context.Context, command string, args ...string) (string, error) {
	cmdMgr := cmd.NewCommandMgr(cmd.WithContext(ctx), cmd.WithTimeout(10*time.Second))
	return cmdMgr.RunWithStdout(command, args...)
}

func parseAMDStatic(data string) (*Info, error) {
	rows, err := decodeAMDRows(data)
	if err != nil {
		return nil, err
	}

	info := &Info{Type: "amd"}
	for _, row := range rows {
		index, ok := amdGPUIndex(row)
		if !ok {
			continue
		}
		device := Device{
			Type:             "amd",
			UUID:             amdStringAt(row, "uuid", "asic.uuid"),
			DriverVersion:    amdStringAt(row, "driver.version", "driver_version", "amdgpu_version"),
			ProcessStatus:    "unavailable",
			Index:            index,
			ProductName:      amdStringAt(row, "asic.market_name", "market_name", "gpu_name"),
			PersistenceMode:  "N/A",
			BusID:            amdStringAt(row, "bus.bdf", "bdf"),
			DisplayActive:    "N/A",
			ECC:              "N/A",
			FanSpeed:         "N/A",
			Temperature:      "N/A",
			PerformanceState: "N/A",
			PowerDraw:        "N/A",
			MaxPowerLimit: amdMetricAt(row, "W",
				"limit.ppt0.max_power_limit",
				"limit.max_power_limit",
				"limit.max_power",
			),
			PowerLimit:  amdMetricAt(row, "W", "limit.ppt0.socket_power_limit", "limit.socket_power_limit"),
			MemUsed:     "N/A",
			MemTotal:    "N/A",
			GPUUtil:     "N/A",
			ComputeMode: "N/A",
			MigMode:     "N/A",
		}
		if device.ProductName == "" {
			device.ProductName = "AMD GPU"
		}
		if device.MaxPowerLimit == "" {
			device.MaxPowerLimit = "N/A"
		}
		if info.DriverVersion == "" {
			info.DriverVersion = amdStringAt(row, "driver.version", "driver_version", "amdgpu_version")
		}
		info.Devices = append(info.Devices, device)
	}
	sort.Slice(info.Devices, func(i, j int) bool {
		return info.Devices[i].Index < info.Devices[j].Index
	})
	return info, nil
}

func applyAMDMetrics(info *Info, data string) error {
	rows, err := decodeAMDRows(data)
	if err != nil {
		return err
	}
	devices := amdDevicesByIndex(info)
	for _, row := range rows {
		index, ok := amdGPUIndex(row)
		if !ok {
			continue
		}
		device, ok := devices[index]
		if !ok {
			continue
		}
		setAMDMetric(&device.GPUUtil, row, "%", "usage.gfx_activity", "usage.gfx", "gfx_activity", "gfx_usage")
		setAMDMetric(&device.Temperature, row, "°C", "temperature.edge", "gpu_temperature", "gpu_temp")
		setAMDMetric(&device.PowerDraw, row, "W", "power.socket_power", "socket_power", "power_usage")
		setAMDMetric(&device.MemUsed, row, "MB", "mem_usage.used_vram", "vram.used", "used_vram", "vram_used")
		setAMDMetric(&device.MemTotal, row, "MB", "mem_usage.total_vram", "vram.total", "total_vram", "vram_total")

		setAMDMetric(&device.FanSpeed, row, "%", "fan.usage")
		if amdMetricAt(row, "%", "fan.usage") == "" {
			speed, speedErr := strconv.ParseFloat(amdStringAt(row, "fan.speed"), 64)
			maximum, maxErr := strconv.ParseFloat(amdStringAt(row, "fan.max"), 64)
			if speedErr == nil && maxErr == nil && !math.IsNaN(speed) && !math.IsInf(speed, 0) && maximum > 0 && !math.IsInf(maximum, 0) && speed >= 0 && speed <= maximum {
				device.FanSpeed = fmt.Sprintf("%.2f %%", speed/maximum*100)
			}
		}
		setAMDMetric(&device.FanRPM, row, "RPM", "fan.rpm")
		setAMDMetric(&device.HotspotTemperature, row, "°C", "temperature.hotspot", "hotspot_temperature")
		setAMDMetric(&device.MemoryTemperature, row, "°C", "temperature.mem")
		setAMDMetric(&device.MemoryActivity, row, "%", "usage.umc_activity")
		setAMDMetric(&device.MediaUtil, row, "%", "usage.mm_activity")
		setAMDMetric(&device.Frequency, row, "MHz", "clock.gfx_0.clk")
		setAMDMetric(&device.MemoryFrequency, row, "MHz", "clock.mem_0.clk")
		correctable := amdStringAt(row, "ecc.total_correctable_count", "ecc.correctable_count")
		uncorrectable := amdStringAt(row, "ecc.total_uncorrectable_count", "ecc.uncorrectable_count")
		if correctable != "" || uncorrectable != "" {
			device.ECCErrors = []ECCError{{Scope: "Total", Correctable: correctable, Uncorrectable: uncorrectable}}
		}
		if value := amdStringAt(row, "perf_level", "performance_level"); value != "" {
			device.PerformanceState = value
		}
	}
	return nil
}

func applyAMDProcesses(info *Info, data string) error {
	rows, err := decodeAMDRows(data)
	if err != nil {
		return err
	}
	devices := amdDevicesByIndex(info)
	for _, row := range rows {
		index, ok := amdGPUIndex(row)
		if !ok {
			continue
		}
		device, ok := devices[index]
		if !ok {
			continue
		}
		device.ProcessStatus = "ok"
		processList, _ := amdValueAt(row, "process_list")
		items := amdObjectList(processList)
		if len(items) == 0 {
			if _, ok := amdValueAt(row, "process_info"); ok {
				items = amdSMIRows{row}
			}
		}
		for _, processRow := range items {
			pid := amdStringAt(processRow, "process_info.pid", "pid")
			if pid == "" || strings.EqualFold(pid, "N/A") {
				continue
			}
			device.Processes = append(device.Processes, Process{
				PID:         pid,
				Type:        "C",
				ProcessName: amdStringAt(processRow, "process_info.name", "name"),
				UsedMemory: amdMetricAt(processRow, "B",
					"process_info.mem_usage",
					"process_info.mem",
					"process_info.memory_usage.vram_mem",
					"process_info.vram_mem",
					"mem_usage",
				),
			})
		}
	}
	return nil
}

func amdDevicesByIndex(info *Info) map[uint]*Device {
	devices := make(map[uint]*Device, len(info.Devices))
	for index := range info.Devices {
		devices[info.Devices[index].Index] = &info.Devices[index]
	}
	return devices
}

func setAMDMetric(target *string, row amdSMIRow, defaultUnit string, paths ...string) {
	if value := amdMetricAt(row, defaultUnit, paths...); value != "" {
		*target = value
	}
}