main.go

package main

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/snappy"
	"github.com/hashicorp/go-retryablehttp"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/prometheus/prompb"
)

// RemoteWriteConfig 远程写入配置
type RemoteWriteConfig struct {
	URL               string
	Username          string
	Password          string
	MaxSamplesPerSend int
}

var config = RemoteWriteConfig{
	URL:               "http://localhost:5080/api/default/prometheus/api/v1/write",
	Username:          "kooksee@163.com",
	Password:          "y3vG9gUDxESqsDMC",
	MaxSamplesPerSend: 10000,
}

const (
	scrapeInterval    = 15 * time.Second
	metricsServerAddr = ":9090"
)

func main() {
	// 启动 HTTP 服务器提供 /metrics 端点
	go func() {
		http.Handle("/metrics", promhttp.Handler())
		log.Printf("Metrics server starting on %s/metrics", metricsServerAddr)
		if err := http.ListenAndServe(metricsServerAddr, nil); err != nil {
			log.Fatalf("Failed to start metrics server: %v", err)
		}
	}()

	// 给服务器一点时间启动
	time.Sleep(100 * time.Millisecond)

	client := makeHttpClient()

	// 主循环：定时收集指标并发送到远程端点
	ticker := time.NewTicker(scrapeInterval)
	defer ticker.Stop()

	// 立即执行一次
	if err := scrapeAndSend(client); err != nil {
		log.Printf("Initial scrape failed: %v", err)
	}

	for range ticker.C {
		if err := scrapeAndSend(client); err != nil {
			log.Printf("Failed to scrape and send metrics: %v", err)
		}
	}
}

// setBasicAuth 设置 HTTP Basic 认证头
func setBasicAuth(req *retryablehttp.Request, username, password string) {
	auth := username + ":" + password
	encoded := base64.StdEncoding.EncodeToString([]byte(auth))
	req.Header.Set("Authorization", "Basic "+encoded)
}

// scrapeAndSend 收集指标并发送到远程端点
func scrapeAndSend(client *retryablehttp.Client) error {
	// 获取当前时间戳，确保同一批次的所有指标使用相同的时间戳
	timestamp := time.Now().UnixMilli()

	// 收集指标
	metricFamilies, err := prometheus.DefaultGatherer.Gather()
	if err != nil {
		return fmt.Errorf("failed to gather metrics: %w", err)
	}

	// 转换为 remote_write 格式
	writeRequest := MetricFamiliesToProtoWriteRequest(metricFamilies, timestamp)
	if len(writeRequest.Timeseries) == 0 {
		log.Println("No timeseries to send")
		return nil
	}

	// 如果指标数量超过 max_samples_per_send，进行分批发送
	timeseries := writeRequest.Timeseries
	totalSamples := countSamples(timeseries)

	if totalSamples > config.MaxSamplesPerSend {
		return sendBatched(client, timeseries, timestamp)
	}

	// 序列化 protobuf
	writeReq := &prompb.WriteRequest{Timeseries: timeseries}
	data, err := proto.Marshal(writeReq)
	if err != nil {
		return fmt.Errorf("failed to marshal write request: %w", err)
	}

	return sendRequest(client, data, len(timeseries), totalSamples)
}

// countSamples 统计 timeseries 中的样本总数
func countSamples(timeseries []prompb.TimeSeries) int {
	total := 0
	for _, ts := range timeseries {
		total += len(ts.Samples)
	}
	return total
}

// sendBatched 分批发送指标（当样本数超过 max_samples_per_send 时）
func sendBatched(client *retryablehttp.Client, timeseries []prompb.TimeSeries, timestamp int64) error {
	totalSent := 0
	batchSize := config.MaxSamplesPerSend

	for i := 0; i < len(timeseries); {
		batch := []prompb.TimeSeries{}
		samplesInBatch := 0

		// 构建一个批次
		j := i
		for j < len(timeseries) && samplesInBatch < batchSize {
			ts := timeseries[j]
			tsSamples := len(ts.Samples)
			// 如果添加这个 timeseries 会超过批次大小，且批次不为空，则停止添加
			if samplesInBatch+tsSamples > batchSize && len(batch) > 0 {
				break
			}
			samplesInBatch += tsSamples
			batch = append(batch, ts)
			j++
		}

		i = j // 更新外层循环索引

		if len(batch) == 0 {
			break
		}

		// 发送批次
		writeReq := &prompb.WriteRequest{Timeseries: batch}
		data, err := proto.Marshal(writeReq)
		if err != nil {
			return fmt.Errorf("failed to marshal batch: %w", err)
		}

		batchSamples := countSamples(batch)
		if err := sendRequest(client, data, len(batch), batchSamples); err != nil {
			return fmt.Errorf("failed to send batch: %w", err)
		}

		totalSent += batchSamples
		log.Printf("Sent batch: %d timeseries, %d samples (total: %d)", len(batch), batchSamples, totalSent)
	}

	return nil
}

// sendRequest 发送 HTTP 请求到远程端点
func sendRequest(client *retryablehttp.Client, data []byte, timeseriesCount, sampleCount int) error {
	compressed := snappy.Encode(nil, data)
	body := bytes.NewReader(compressed)

	// 创建 HTTP 请求
	req, err := retryablehttp.NewRequest("POST", config.URL, body)
	if err != nil {
		return fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-protobuf")
	req.Header.Set("Content-Encoding", "snappy")
	setBasicAuth(req, config.Username, config.Password)

	// 发送请求
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("failed to send request: %w", err)
	}
	defer resp.Body.Close()

	// 读取响应
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("failed to read response body: %w", err)
	}

	// 检查响应状态
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("remote write failed with status %d: %s", resp.StatusCode, string(respBody))
	}

	log.Printf("Successfully sent %d timeseries (%d samples) to OpenObserve", timeseriesCount, sampleCount)
	return nil
}

// OpenObserve remote_write 配置参考：
// remote_write:
//   - url: http://localhost:5080/api/default/prometheus/api/v1/write
//     queue_config:
//       max_samples_per_send: 10000
//     basic_auth:
//       username: kooksee@163.com
//       password: y3vG9gUDxESqsDMC

// MetricFamiliesToProtoWriteRequest 将 MetricFamily 列表转换为 Prometheus remote_write 的 WriteRequest
// timestamp 参数确保同一批次的所有指标使用相同的时间戳
func MetricFamiliesToProtoWriteRequest(metricFamilies []*dto.MetricFamily, timestamp int64) *prompb.WriteRequest {
	writeRequest := &prompb.WriteRequest{
		Timeseries: make([]prompb.TimeSeries, 0),
	}

	for _, family := range metricFamilies {
		switch family.GetType() {
		case dto.MetricType_COUNTER:
			appendCounterToWriteRequest(family, writeRequest, timestamp)
		case dto.MetricType_GAUGE:
			appendGaugeToWriteRequest(family, writeRequest, timestamp)
		case dto.MetricType_HISTOGRAM:
			appendHistogramToWriteRequest(family, writeRequest, timestamp)
		case dto.MetricType_SUMMARY:
			appendSummaryToWriteRequest(family, writeRequest, timestamp)
		default:
			log.Printf("Unsupported metric type: %v for metric %s", family.GetType(), family.GetName())
		}
	}

	return writeRequest
}

func ExtractLabels(metric *dto.Metric) []prompb.Label {
	labels := make([]prompb.Label, 0, len(metric.Label))

	for _, pair := range metric.GetLabel() {
		if *pair.Name != "" && *pair.Value != "" {
			labels = append(labels, prompb.Label{
				Name:  *pair.Name,
				Value: *pair.Value,
			})
		}
	}
	return labels
}

func appendCounterToWriteRequest(family *dto.MetricFamily, wr *prompb.WriteRequest, timestamp int64) {
	for _, metric := range family.GetMetric() {
		labels := ExtractLabels(metric)
		labels = append(labels, prompb.Label{
			Name:  "__name__",
			Value: family.GetName(),
		})

		samples := []prompb.Sample{
			{
				Value:     metric.GetCounter().GetValue(),
				Timestamp: timestamp,
			},
		}

		wr.Timeseries = append(wr.Timeseries, prompb.TimeSeries{
			Labels:  labels,
			Samples: samples,
		})
	}
}

// appendGaugeToWriteRequest 将 GAUGE 类型的指标添加到 WriteRequest
func appendGaugeToWriteRequest(family *dto.MetricFamily, wr *prompb.WriteRequest, timestamp int64) {
	for _, metric := range family.GetMetric() {
		labels := ExtractLabels(metric)
		labels = append(labels, prompb.Label{
			Name:  "__name__",
			Value: family.GetName(),
		})

		samples := []prompb.Sample{
			{
				Value:     metric.GetGauge().GetValue(),
				Timestamp: timestamp,
			},
		}

		wr.Timeseries = append(wr.Timeseries, prompb.TimeSeries{
			Labels:  labels,
			Samples: samples,
		})
	}
}

func appendHistogramToWriteRequest(family *dto.MetricFamily, wr *prompb.WriteRequest, timestamp int64) {
	for _, metric := range family.GetMetric() {
		labels := ExtractLabels(metric)
		hist := metric.GetHistogram()

		// hist count
		countLabels := make([]prompb.Label, len(labels), len(labels)+1)
		copy(countLabels, labels)
		countLabels = append(countLabels, prompb.Label{
			Name:  "__name__",
			Value: fmt.Sprintf("%s_count", family.GetName()),
		})
		wr.Timeseries = append(wr.Timeseries, prompb.TimeSeries{
			Labels: countLabels,
			Samples: []prompb.Sample{
				{
					Value:     float64(hist.GetSampleCount()),
					Timestamp: timestamp,
				},
			},
		})

		// hist sum
		sumLabels := make([]prompb.Label, len(labels), len(labels)+1)
		copy(sumLabels, labels)
		sumLabels = append(sumLabels, prompb.Label{
			Name:  "__name__",
			Value: fmt.Sprintf("%s_sum", family.GetName()),
		})
		wr.Timeseries = append(wr.Timeseries, prompb.TimeSeries{
			Labels: sumLabels,
			Samples: []prompb.Sample{
				{
					Value:     hist.GetSampleSum(),
					Timestamp: timestamp,
				},
			},
		})

		// hist bucket
		for _, bucket := range hist.GetBucket() {
			bucketLabels := make([]prompb.Label, len(labels), len(labels)+2)
			copy(bucketLabels, labels)
			bucketLabels = append(bucketLabels, prompb.Label{
				Name:  "le",
				Value: fmt.Sprintf("%g", bucket.GetUpperBound()),
			})
			bucketLabels = append(bucketLabels, prompb.Label{
				Name:  "__name__",
				Value: fmt.Sprintf("%s_bucket", family.GetName()),
			})

			wr.Timeseries = append(wr.Timeseries, prompb.TimeSeries{
				Labels: bucketLabels,
				Samples: []prompb.Sample{
					{
						Value:     float64(bucket.GetCumulativeCount()),
						Timestamp: timestamp,
					},
				},
			})
		}
	}
}

// appendSummaryToWriteRequest 将 SUMMARY 类型的指标添加到 WriteRequest
func appendSummaryToWriteRequest(family *dto.MetricFamily, wr *prompb.WriteRequest, timestamp int64) {
	for _, metric := range family.GetMetric() {
		labels := ExtractLabels(metric)
		summary := metric.GetSummary()

		// summary count
		countLabels := make([]prompb.Label, len(labels), len(labels)+1)
		copy(countLabels, labels)
		countLabels = append(countLabels, prompb.Label{
			Name:  "__name__",
			Value: fmt.Sprintf("%s_count", family.GetName()),
		})
		wr.Timeseries = append(wr.Timeseries, prompb.TimeSeries{
			Labels: countLabels,
			Samples: []prompb.Sample{
				{
					Value:     float64(summary.GetSampleCount()),
					Timestamp: timestamp,
				},
			},
		})

		// summary sum
		sumLabels := make([]prompb.Label, len(labels), len(labels)+1)
		copy(sumLabels, labels)
		sumLabels = append(sumLabels, prompb.Label{
			Name:  "__name__",
			Value: fmt.Sprintf("%s_sum", family.GetName()),
		})
		wr.Timeseries = append(wr.Timeseries, prompb.TimeSeries{
			Labels: sumLabels,
			Samples: []prompb.Sample{
				{
					Value:     summary.GetSampleSum(),
					Timestamp: timestamp,
				},
			},
		})

		// summary quantiles
		for _, quantile := range summary.GetQuantile() {
			quantileLabels := make([]prompb.Label, len(labels), len(labels)+2)
			copy(quantileLabels, labels)
			quantileLabels = append(quantileLabels, prompb.Label{
				Name:  "quantile",
				Value: fmt.Sprintf("%g", quantile.GetQuantile()),
			})
			quantileLabels = append(quantileLabels, prompb.Label{
				Name:  "__name__",
				Value: family.GetName(),
			})

			wr.Timeseries = append(wr.Timeseries, prompb.TimeSeries{
				Labels: quantileLabels,
				Samples: []prompb.Sample{
					{
						Value:     quantile.GetValue(),
						Timestamp: timestamp,
					},
				},
			})
		}
	}
}

// ScrapeTarget scrapes an http endpoint for prometheus metrics
func ScrapeTarget(targetURL string) (string, map[string]*dto.MetricFamily, error) {
	resp, err := retryablehttp.Get(targetURL)
	if err != nil {
		return "", nil, fmt.Errorf("error fetching metrics from target: %w", err)
	}
	//nolint:staticcheck // SA5001 Ignore error here
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", nil, fmt.Errorf("error reading response body: %w", err)
	}

	parser := expfmt.TextParser{}

	metricFamily, err := parser.TextToMetricFamilies(bytes.NewReader(body))
	return string(body), metricFamily, err
}

func makeHttpClient() *retryablehttp.Client {
	client := retryablehttp.NewClient()
	client.RetryMax = 10
	return client
}