package collector

import (
	"encoding/json"
	"os/exec"
	"strings"
)

type GPUInfo struct {
	Usage       float64 `json:"usage"`
	Temperature float64 `json:"temperature"`
	MemoryTotal uint64  `json:"memory_total"`
	MemoryUsed  uint64  `json:"memory_used"`
	Name        string  `json:"name"`
}

type nvidiaSMIOutput struct {
	GPUs []struct {
		Name         string `json:"name"`
		Utilization  struct {
			GPU string `json:"gpu"`
		} `json:"utilization"`
		Temperature struct {
			GPU string `json:"gpu_temp"`
		} `json:"temperature"`
		Memory struct {
			Total string `json:"total"`
			Used  string `json:"used"`
		} `json:"fb_memory_usage"`
	} `json:"gpus"`
}

func CollectGPU() (*GPUInfo, error) {
	// Try NVIDIA GPU first
	info, err := collectNvidiaGPU()
	if err == nil && info != nil {
		return info, nil
	}

	// Return empty GPU info if no GPU detected
	return &GPUInfo{
		Name: "No GPU detected",
	}, nil
}

func collectNvidiaGPU() (*GPUInfo, error) {
	cmd := exec.Command("nvidia-smi", "--query-gpu=name,utilization.gpu,temperature.gpu,memory.total,memory.used", "--format=csv,noheader,nounits")
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
	if len(lines) == 0 {
		return nil, nil
	}

	fields := strings.Split(lines[0], ",")
	if len(fields) < 5 {
		return nil, nil
	}

	info := &GPUInfo{
		Name: strings.TrimSpace(fields[0]),
	}

	// Parse values (ignore errors, use defaults)
	var usage, temp, memTotal, memUsed float64
	json.Unmarshal([]byte(strings.TrimSpace(fields[1])), &usage)
	json.Unmarshal([]byte(strings.TrimSpace(fields[2])), &temp)
	json.Unmarshal([]byte(strings.TrimSpace(fields[3])), &memTotal)
	json.Unmarshal([]byte(strings.TrimSpace(fields[4])), &memUsed)

	info.Usage = usage
	info.Temperature = temp
	info.MemoryTotal = uint64(memTotal) * 1024 * 1024 // Convert MB to bytes
	info.MemoryUsed = uint64(memUsed) * 1024 * 1024

	return info, nil
}