package server

import (
	"context"
	"crypto/subtle"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"vdi-cloud/platform/pkg/proxmox"
)

// gpuAgentSample is the small, vendor-neutral contract implemented by a node
// collector. The collector may use nvidia-smi, rocm-smi, intel_gpu_top or a
// vendor API, while the platform keeps the same health model for every GPU.
type gpuAgentSample struct {
	Address            string   `json:"address" binding:"required"`
	ObservedAt         string   `json:"observedAt"`
	TemperatureC       *float64 `json:"temperatureC"`
	UtilizationPercent *float64 `json:"utilizationPercent"`
	MemoryUsedBytes    *int64   `json:"memoryUsedBytes"`
	MemoryTotalBytes   *int64   `json:"memoryTotalBytes"`
	Driver             string   `json:"driver"`
	DriverBinding      string   `json:"driverBinding"`
}

type gpuAgentReport struct {
	Node    string           `json:"node" binding:"required,max=100"`
	Samples []gpuAgentSample `json:"samples" binding:"required,min=1,max=256"`
}

func (s *Server) gpuAgentAuth(c *gin.Context) {
	if strings.TrimSpace(s.Config.GPUAgentToken) == "" {
		s.fail(c, 503, "gpu_agent_unconfigured", "尚未配置 GPU 节点采集令牌")
		return
	}
	provided := c.GetHeader("X-GPU-Agent-Token")
	if len(provided) != len(s.Config.GPUAgentToken) || subtle.ConstantTimeCompare([]byte(provided), []byte(s.Config.GPUAgentToken)) != 1 {
		s.fail(c, 401, "gpu_agent_unauthorized", "GPU 节点采集令牌无效")
		return
	}
	c.Next()
}

func validGPUHealth(sample gpuAgentSample) bool {
	if sample.TemperatureC != nil && (*sample.TemperatureC < -40 || *sample.TemperatureC > 150) {
		return false
	}
	if sample.UtilizationPercent != nil && (*sample.UtilizationPercent < 0 || *sample.UtilizationPercent > 100) {
		return false
	}
	if sample.MemoryUsedBytes != nil && *sample.MemoryUsedBytes < 0 {
		return false
	}
	if sample.MemoryTotalBytes != nil && *sample.MemoryTotalBytes < 0 {
		return false
	}
	if sample.MemoryUsedBytes != nil && sample.MemoryTotalBytes != nil && *sample.MemoryUsedBytes > *sample.MemoryTotalBytes {
		return false
	}
	_, ok := proxmox.NormalizePCIAddress(sample.Address)
	return ok
}

func (s *Server) ingestGPUHealth(c *gin.Context) {
	var in gpuAgentReport
	if !s.decode(c, &in) {
		return
	}
	in.Node = strings.TrimSpace(in.Node)
	if in.Node == "" || strings.ContainsAny(in.Node, "/\\?# ") {
		s.fail(c, 400, "invalid_node", "节点名称无效")
		return
	}
	ctx := c.Request.Context()
	tx, err := s.DB.BeginTxx(ctx, nil)
	if err != nil {
		s.dbError(c, err)
		return
	}
	defer tx.Rollback()
	for _, sample := range in.Samples {
		if !validGPUHealth(sample) {
			s.fail(c, 400, "invalid_gpu_health", "GPU 健康数据超出范围或 PCI 地址无效")
			return
		}
		address, _ := proxmox.NormalizePCIAddress(sample.Address)
		observedAt := time.Now().UTC()
		if strings.TrimSpace(sample.ObservedAt) != "" {
			observedAt, err = time.Parse(time.RFC3339Nano, sample.ObservedAt)
			if err != nil || observedAt.After(time.Now().Add(2*time.Minute)) || observedAt.Before(time.Now().Add(-10*time.Minute)) {
				s.fail(c, 400, "invalid_gpu_health", "采集时间必须是最近 10 分钟内的 RFC3339 时间")
				return
			}
			observedAt = observedAt.UTC()
		}
		_, err = tx.ExecContext(ctx, `INSERT INTO gpu_health_samples(node,address,observed_at,temperature_c,utilization_percent,memory_used_bytes,memory_total_bytes,driver,driver_binding,source)
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,'node-agent') ON CONFLICT DO NOTHING`, in.Node, address, observedAt, sample.TemperatureC, sample.UtilizationPercent, sample.MemoryUsedBytes, sample.MemoryTotalBytes, strings.TrimSpace(sample.Driver), strings.TrimSpace(sample.DriverBinding))
		if err != nil {
			s.dbError(c, err)
			return
		}
	}
	if _, err = tx.ExecContext(ctx, "DELETE FROM gpu_health_samples WHERE observed_at < now() - interval '7 days'"); err != nil {
		s.dbError(c, err)
		return
	}
	if err = tx.Commit(); err != nil {
		s.dbError(c, err)
		return
	}
	c.JSON(202, gin.H{"accepted": len(in.Samples), "node": in.Node})
}

type latestGPUHealth struct {
	Address            string    `db:"address"`
	ObservedAt         time.Time `db:"observed_at"`
	TemperatureC       *float64  `db:"temperature_c"`
	UtilizationPercent *float64  `db:"utilization_percent"`
	MemoryUsedBytes    *int64    `db:"memory_used_bytes"`
	MemoryTotalBytes   *int64    `db:"memory_total_bytes"`
	Driver             string    `db:"driver"`
	DriverBinding      string    `db:"driver_binding"`
	Source             string    `db:"source"`
}

func (s *Server) latestGPUHealth(ctx context.Context, node string) (map[string]latestGPUHealth, error) {
	// Kept separate from inventory assembly so a missing collector is an
	// ordinary empty result and remains visibly "not_collected" in the API.
	items := []latestGPUHealth{}
	query := `SELECT DISTINCT ON(address) address,observed_at,temperature_c,utilization_percent,memory_used_bytes,memory_total_bytes,driver,driver_binding,source
FROM gpu_health_samples WHERE node=$1 AND observed_at > now()-interval '10 minutes' ORDER BY address,observed_at DESC`
	if err := s.DB.SelectContext(ctx, &items, query, node); err != nil {
		return nil, err
	}
	result := make(map[string]latestGPUHealth, len(items))
	for _, item := range items {
		result[item.Address] = item
	}
	return result, nil
}