hi-server/pkg/loki/loki.go

package loki

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"time"
)

// Client Loki 客户端
type Client struct {
	url        string
	httpClient *http.Client
}

// NewClient 创建新的 Loki 客户端
// url: Loki 服务地址，例如 http://154.12.35.103:3100
func NewClient(url string) *Client {
	return &Client{
		url: url,
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
		},
	}
}

// InviteCodeStats 邀请码统计数据
type InviteCodeStats struct {
	MacClicks        int64 `json:"mac_clicks"`         // Mac 下载点击数
	WindowsClicks    int64 `json:"windows_clicks"`     // Windows 下载点击数
	LastMonthMac     int64 `json:"last_month_mac"`     // 上月 Mac 下载数
	LastMonthWindows int64 `json:"last_month_windows"` // 上月 Windows 下载数
}

// LokiQueryResponse Loki 查询响应结构
type LokiQueryResponse struct {
	Status string `json:"status"`
	Data   struct {
		ResultType string `json:"resultType"`
		Result     []struct {
			Stream map[string]string `json:"stream"`
			Values [][]string        `json:"values"` // [[timestamp, log_line], ...]
		} `json:"result"`
	} `json:"data"`
}

// GetInviteCodeStats 获取指定邀请码的下载统计
// inviteCode: 邀请码
// days: 统计天数（默认30天）
func (c *Client) GetInviteCodeStats(ctx context.Context, inviteCode string, days int) (*InviteCodeStats, error) {
	if days <= 0 {
		days = 30
	}

	now := time.Now().UTC()
	startTime := now.Add(-time.Duration(days) * 24 * time.Hour)

	// 上月时间范围
	lastMonthEnd := startTime
	lastMonthStart := startTime.Add(-time.Duration(days) * 24 * time.Hour)

	// 查询本月数据
	thisMonthStats, err := c.queryPeriodStats(ctx, inviteCode, startTime, now)
	if err != nil {
		return nil, fmt.Errorf("查询本月数据失败: %w", err)
	}

	// 查询上月数据
	lastMonthStats, err := c.queryPeriodStats(ctx, inviteCode, lastMonthStart, lastMonthEnd)
	if err != nil {
		return nil, fmt.Errorf("查询上月数据失败: %w", err)
	}

	return &InviteCodeStats{
		MacClicks:        thisMonthStats.MacClicks,
		WindowsClicks:    thisMonthStats.WindowsClicks,
		LastMonthMac:     lastMonthStats.MacClicks,
		LastMonthWindows: lastMonthStats.WindowsClicks,
	}, nil
}

// GetInviteCodeMonthlyStats 获取指定邀请码在自然月维度的下载统计。
// currentTime 所在月份作为本月，上一自然月作为上月。
func (c *Client) GetInviteCodeMonthlyStats(ctx context.Context, inviteCode string, currentTime time.Time) (*InviteCodeStats, error) {
	if currentTime.IsZero() {
		currentTime = time.Now()
	}

	currentMonthStart := time.Date(currentTime.Year(), currentTime.Month(), 1, 0, 0, 0, 0, currentTime.Location())
	lastMonthStart := currentMonthStart.AddDate(0, -1, 0)

	thisMonthStats, err := c.queryPeriodStats(ctx, inviteCode, currentMonthStart, currentTime)
	if err != nil {
		return nil, fmt.Errorf("查询本月数据失败: %w", err)
	}

	lastMonthStats, err := c.queryPeriodStats(ctx, inviteCode, lastMonthStart, currentMonthStart)
	if err != nil {
		return nil, fmt.Errorf("查询上月数据失败: %w", err)
	}

	return &InviteCodeStats{
		MacClicks:        thisMonthStats.MacClicks,
		WindowsClicks:    thisMonthStats.WindowsClicks,
		LastMonthMac:     lastMonthStats.MacClicks,
		LastMonthWindows: lastMonthStats.WindowsClicks,
	}, nil
}

// queryPeriodStats 查询指定时间范围的统计数据
func (c *Client) queryPeriodStats(ctx context.Context, inviteCode string, startTime, endTime time.Time) (*InviteCodeStats, error) {
	// 构建 Loki 查询
	query := fmt.Sprintf(`{job="nginx_access", invite_code="%s"}`, inviteCode)

	apiURL := fmt.Sprintf("%s/loki/api/v1/query_range", c.url)

	params := url.Values{}
	params.Add("query", query)
	params.Add("start", startTime.Format(time.RFC3339))
	params.Add("end", endTime.Format(time.RFC3339))
	params.Add("limit", "5000")

	fullURL := fmt.Sprintf("%s?%s", apiURL, params.Encode())

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %w", err)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("发送请求失败: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("Loki 返回错误状态码 %d: %s", resp.StatusCode, string(body))
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应失败: %w", err)
	}

	var lokiResp LokiQueryResponse
	if err := json.Unmarshal(body, &lokiResp); err != nil {
		return nil, fmt.Errorf("解析响应失败: %w", err)
	}

	// 解析日志行统计 Mac 和 Windows 下载
	stats := &InviteCodeStats{}

	// Nginx combined log format regex
	// 格式: IP - - [time] "METHOD URI VERSION" STATUS BYTES "REFERER" "UA"
	logPattern := regexp.MustCompile(`"[A-Z]+ ([^ ]+) `)

	for _, result := range lokiResp.Data.Result {
		for _, value := range result.Values {
			if len(value) < 2 {
				continue
			}
			logLine := value[1]

			// 提取 URI
			matches := logPattern.FindStringSubmatch(logLine)
			if len(matches) < 2 {
				continue
			}
			uri := strings.ToLower(matches[1])

			// 统计平台下载
			if strings.Contains(uri, "mac") {
				stats.MacClicks++
			} else if strings.Contains(uri, "windows") {
				stats.WindowsClicks++
			}
		}
	}

	return stats, nil
}