aiops/backend/controller/prometheus_proxy.go

package controller

import (
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"ezops/model"

	"github.com/gin-gonic/gin"
)

// PrometheusProxyHandler Prometheus API 代理 — 转发 PromQL 查询到 Prometheus 服务器
type PrometheusProxyHandler struct{}

func NewPrometheusProxyHandler() *PrometheusProxyHandler {
	return &PrometheusProxyHandler{}
}

// PromConnectReq 连接 Prometheus 请求
type PromConnectReq struct {
	URL string `json:"url" binding:"required"`
}

// PromConnect 连接 Prometheus 并返回概览数据
// POST /api/v1/prometheus/connect
func (h *PrometheusProxyHandler) Connect(c *gin.Context) {
	var req PromConnectReq
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "参数错误: " + err.Error()})
		return
	}

	client := &http.Client{Timeout: 15 * time.Second}

	// 测试连接：查询 up 指标
	status, targetCount := h.testConnect(client, req.URL)
	if !status {
		c.JSON(http.StatusBadGateway, model.Response{Code: 502, Message: "无法连接到 Prometheus，请检查 URL"})
		return
	}

	c.JSON(http.StatusOK, model.Response{
		Code:    200,
		Message: "success",
		Data: gin.H{
			"status":  "connected",
			"targets": targetCount,
		},
	})
}

// PromQuery 执行 PromQL 即时查询
// POST /api/v1/prometheus/query
func (h *PrometheusProxyHandler) Query(c *gin.Context) {
	var req struct {
		URL   string `json:"url" binding:"required"`
		Query string `json:"query" binding:"required"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "参数错误"})
		return
	}

	client := &http.Client{Timeout: 15 * time.Second}
	result, err := h.execPromQL(client, req.URL, req.Query)
	if err != nil {
		c.JSON(http.StatusBadGateway, model.Response{Code: 502, Message: "查询失败: " + err.Error()})
		return
	}

	c.JSON(http.StatusOK, model.Response{Code: 200, Message: "success", Data: result})
}

// PromQueryRange 执行 PromQL 范围查询
// POST /api/v1/prometheus/query_range
func (h *PrometheusProxyHandler) QueryRange(c *gin.Context) {
	var req struct {
		URL   string `json:"url" binding:"required"`
		Query string `json:"query" binding:"required"`
		Start string `json:"start"` // RFC3339
		End   string `json:"end"`
		Step  string `json:"step"` // e.g. "15s", "1m"
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "参数错误"})
		return
	}
	if req.Step == "" {
		req.Step = "1m"
	}

	client := &http.Client{Timeout: 30 * time.Second}

	u := fmt.Sprintf("%s/api/v1/query_range?query=%s", strings.TrimRight(req.URL, "/"),
		url.QueryEscape(req.Query))
	if req.Start != "" {
		u += "&start=" + url.QueryEscape(req.Start)
	}
	if req.End != "" {
		u += "&end=" + url.QueryEscape(req.End)
	}
	u += "&step=" + url.QueryEscape(req.Step)

	resp, err := client.Get(u)
	if err != nil {
		c.JSON(http.StatusBadGateway, model.Response{Code: 502, Message: "查询失败: " + err.Error()})
		return
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	var promResult map[string]interface{}
	json.Unmarshal(body, &promResult)

	c.JSON(http.StatusOK, model.Response{Code: 200, Message: "success", Data: promResult})
}

// PromDashboard 获取预定义仪表盘数据（K8s / Node / Pod / JVM）
// POST /api/v1/prometheus/dashboard
func (h *PrometheusProxyHandler) Dashboard(c *gin.Context) {
	var req struct {
		URL  string `json:"url" binding:"required"`
		Type string `json:"type" binding:"required"` // k8s, node, pod, jvm
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "参数错误"})
		return
	}

	client := &http.Client{Timeout: 30 * time.Second}
	var result map[string]interface{}

	switch req.Type {
	case "node":
		result = h.dashboardNode(client, req.URL)
	case "pod":
		result = h.dashboardPod(client, req.URL)
	case "jvm":
		result = h.dashboardJVM(client, req.URL)
	default:
		result = h.dashboardK8s(client, req.URL)
	}

	c.JSON(http.StatusOK, model.Response{Code: 200, Message: "success", Data: result})
}

// ============================================
// 内部方法
// ============================================

func (h *PrometheusProxyHandler) testConnect(client *http.Client, baseURL string) (bool, int) {
	result, err := h.execPromQL(client, baseURL, "count(up)")
	if err != nil {
		return false, 0
	}
	targets := 0
	if data, ok := result["data"].(map[string]interface{}); ok {
		if res, ok := data["result"].([]interface{}); ok {
			if len(res) > 0 {
				if val, ok := res[0].(map[string]interface{})["value"].([]interface{}); ok && len(val) > 1 {
					targets, _ = strconv.Atoi(fmt.Sprintf("%v", val[1]))
				}
			}
		}
	}
	return true, targets
}

func (h *PrometheusProxyHandler) execPromQL(client *http.Client, baseURL, query string) (map[string]interface{}, error) {
	u := fmt.Sprintf("%s/api/v1/query?query=%s", strings.TrimRight(baseURL, "/"), url.QueryEscape(query))
	resp, err := client.Get(u)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	var result map[string]interface{}
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, err
	}
	if status, ok := result["status"].(string); !ok || status != "success" {
		return nil, fmt.Errorf("prometheus error: %v", result)
	}
	return result, nil
}

func parseVector(result map[string]interface{}) []map[string]interface{} {
	var out []map[string]interface{}
	data, _ := result["data"].(map[string]interface{})
	res, _ := data["result"].([]interface{})
	for _, r := range res {
		rm, _ := r.(map[string]interface{})
		metric, _ := rm["metric"].(map[string]interface{})
		value, _ := rm["value"].([]interface{})
		entry := map[string]interface{}{}
		for k, v := range metric {
			entry[k] = v
		}
		if len(value) > 1 {
			if v, err := strconv.ParseFloat(fmt.Sprintf("%v", value[1]), 64); err == nil {
				entry["value"] = math.Round(v*100) / 100
			}
		}
		out = append(out, entry)
	}
	return out
}

// ============================================
// 预定义仪表盘查询
// ============================================

func (h *PrometheusProxyHandler) dashboardNode(client *http.Client, baseURL string) map[string]interface{} {
	results := gin.H{}

	// 节点概览表
	results["nodes"] = parseVector(h.mustQuery(client, baseURL, `node_uname_info - 0`))

	// CPU
	results["cpuUsage"] = parseVector(h.mustQuery(client, baseURL,
		`(1 - avg(irate(node_cpu_seconds_total{mode="idle"}[5m])) by (instance)) * 100`))

	// 内存
	results["memUsage"] = parseVector(h.mustQuery(client, baseURL,
		`(1 - (node_memory_MemAvailable_bytes / node_memory_MemTotal_bytes)) * 100`))

	// 磁盘
	results["diskUsage"] = parseVector(h.mustQuery(client, baseURL,
		`(1 - sum(node_filesystem_avail_bytes{fstype=~"xfs|ext.*"}) by (instance) / sum(node_filesystem_size_bytes{fstype=~"xfs|ext.*"}) by (instance)) * 100`))

	// Load
	results["load5"] = parseVector(h.mustQuery(client, baseURL, `node_load5`))

	// 网络
	results["netRx"] = parseVector(h.mustQuery(client, baseURL,
		`max(irate(node_network_receive_bytes_total[5m])) by (instance)`))
	results["netTx"] = parseVector(h.mustQuery(client, baseURL,
		`max(irate(node_network_transmit_bytes_total[5m])) by (instance)`))

	results["up"] = parseVector(h.mustQuery(client, baseURL, `up`))

	return results
}

func (h *PrometheusProxyHandler) dashboardPod(client *http.Client, baseURL string) map[string]interface{} {
	results := gin.H{}

	results["cpuTop"] = parseVector(h.mustQuery(client, baseURL,
		`topk(10, sum(rate(container_cpu_usage_seconds_total{container!="",pod!=""}[5m])) by (namespace, pod))`))

	results["memTop"] = parseVector(h.mustQuery(client, baseURL,
		`topk(10, sum(container_memory_working_set_bytes{container!="",pod!=""}) by (namespace, pod))`))

	results["netRxTop"] = parseVector(h.mustQuery(client, baseURL,
		`topk(10, sum(rate(container_network_receive_bytes_total{pod!=""}[5m])) by (namespace, pod))`))

	results["netTxTop"] = parseVector(h.mustQuery(client, baseURL,
		`topk(10, sum(rate(container_network_transmit_bytes_total{pod!=""}[5m])) by (namespace, pod))`))

	results["restartTop"] = parseVector(h.mustQuery(client, baseURL,
		`topk(10, sum(kube_pod_container_status_restarts_total) by (namespace, pod))`))

	results["running"] = parseVector(h.mustQuery(client, baseURL,
		`sum(kube_pod_status_phase{phase="Running"})`))
	results["pending"] = parseVector(h.mustQuery(client, baseURL,
		`sum(kube_pod_status_phase{phase="Pending"}) or vector(0)`))

	return results
}

func (h *PrometheusProxyHandler) dashboardJVM(client *http.Client, baseURL string) map[string]interface{} {
	results := gin.H{}

	results["uptime"] = parseVector(h.mustQuery(client, baseURL, `process_uptime_seconds`))
	results["heapUsed"] = parseVector(h.mustQuery(client, baseURL,
		`sum(jvm_memory_used_bytes{area="heap"}) / sum(jvm_memory_max_bytes{area="heap"}) * 100`))
	results["nonHeapUsed"] = parseVector(h.mustQuery(client, baseURL,
		`sum(jvm_memory_used_bytes{area="nonheap"}) / sum(jvm_memory_max_bytes{area="nonheap"}) * 100`))
	results["threadCount"] = parseVector(h.mustQuery(client, baseURL,
		`jvm_threads_live_threads`))
	results["gcPause"] = parseVector(h.mustQuery(client, baseURL,
		`rate(jvm_gc_pause_seconds_sum[5m])`))
	results["cpuLoad"] = parseVector(h.mustQuery(client, baseURL,
		`process_cpu_usage`))

	return results
}

func (h *PrometheusProxyHandler) dashboardK8s(client *http.Client, baseURL string) map[string]interface{} {
	results := gin.H{}

	// 节点数
	results["nodeCount"] = parseVector(h.mustQuery(client, baseURL, `count(kube_node_info)`))

	// Pod 总数
	results["podCount"] = parseVector(h.mustQuery(client, baseURL,
		`count(kube_pod_info{created_by_kind!~"<none>|Job"})`))

	// 节点 CPU 利用率
	results["nodeCPU"] = parseVector(h.mustQuery(client, baseURL,
		`sum(irate(container_cpu_usage_seconds_total{container!=""}[2m])) by (node) / sum(kube_pod_container_resource_requests{resource="cpu"}) by (node) * 100`))

	// 节点内存利用率
	results["nodeMem"] = parseVector(h.mustQuery(client, baseURL,
		`sum(container_memory_working_set_bytes{container!=""}) by (node) / sum(kube_pod_container_resource_requests{resource="memory"}) by (node) * 100`))

	// Namespace 资源统计
	results["nsPods"] = parseVector(h.mustQuery(client, baseURL,
		`count(kube_pod_info) by (namespace)`))

	// 部署统计
	results["deployCount"] = parseVector(h.mustQuery(client, baseURL,
		`count(kube_deployment_metadata_generation)`))

	return results
}

func (h *PrometheusProxyHandler) mustQuery(client *http.Client, baseURL, query string) map[string]interface{} {
	r, err := h.execPromQL(client, baseURL, query)
	if err != nil {
		return map[string]interface{}{"data": map[string]interface{}{"result": []interface{}{}}}
	}
	return r
}