gotidb/pkg/api/rest.go

package api

import (
	"context"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"

	"git.pyer.club/kingecg/gotidb/pkg/manager"
	"git.pyer.club/kingecg/gotidb/pkg/model"
)

// RESTServer REST API服务
type RESTServer struct {
	dataManager *manager.DataManager
	router      *gin.Engine
	server      *http.Server
}

// RESTWriteRequest 写入请求
type RESTWriteRequest struct {
	DeviceID   string            `json:"device_id"`
	MetricCode string            `json:"metric_code"`
	Labels     map[string]string `json:"labels"`
	Value      interface{}       `json:"value"`
	Timestamp  *time.Time        `json:"timestamp,omitempty"`
}

type Response map[string]any

// BatchWriteRequest 批量写入请求
type BatchWriteRequest struct {
	Points []RESTWriteRequest `json:"points"`
}

// QueryRequest 查询请求
type QueryRequest struct {
	DeviceID   string                 `json:"device_id"`
	MetricCode string                 `json:"metric_code"`
	Labels     map[string]string      `json:"labels"`
	QueryType  string                 `json:"query_type"`
	Params     map[string]interface{} `json:"params"`
}

// NewRESTServer 创建一个新的REST API服务
func NewRESTServer(dataManager *manager.DataManager) *RESTServer {
	router := gin.Default()

	server := &RESTServer{
		dataManager: dataManager,
		router:      router,
	}

	server.setupRoutes()

	return server
}

// setupRoutes 设置路由
func (s *RESTServer) setupRoutes() {
	// 健康检查
	s.router.GET("/health", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"status": "ok",
		})
	})

	// API版本
	v1 := s.router.Group("/api/v1")
	{
		// 写入数据
		v1.POST("/write", s.handleWrite)

		// 批量写入数据
		v1.POST("/batch_write", s.handleBatchWrite)

		// 查询数据
		v1.POST("/query", s.handleQuery)
	}
}

// handleWrite 处理写入请求
func (s *RESTServer) handleWrite(c *gin.Context) {
	var req RESTWriteRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"error": "Invalid request: " + err.Error(),
		})
		return
	}

	// 创建数据点
	timestamp := time.Now()
	if req.Timestamp != nil {
		timestamp = *req.Timestamp
	}

	id := model.DataPointID{
		DeviceID:   req.DeviceID,
		MetricCode: req.MetricCode,
		Labels:     req.Labels,
	}

	value := model.DataValue{
		Timestamp: timestamp,
		Value:     req.Value,
	}

	// 写入数据
	if err := s.dataManager.Write(c.Request.Context(), id, value); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "Failed to write data: " + err.Error(),
		})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"status": "ok",
	})
}

// handleBatchWrite 处理批量写入请求
func (s *RESTServer) handleBatchWrite(c *gin.Context) {
	var req BatchWriteRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"error": "Invalid request: " + err.Error(),
		})
		return
	}

	// 创建批量写入数据
	batch := make([]struct {
		ID    model.DataPointID
		Value model.DataValue
	}, len(req.Points))

	for i, point := range req.Points {
		timestamp := time.Now()
		if point.Timestamp != nil {
			timestamp = *point.Timestamp
		}

		batch[i].ID = model.DataPointID{
			DeviceID:   point.DeviceID,
			MetricCode: point.MetricCode,
			Labels:     point.Labels,
		}

		batch[i].Value = model.DataValue{
			Timestamp: timestamp,
			Value:     point.Value,
		}
	}

	// 批量写入数据
	// 使用当前时间作为批量写入的时间戳
	if err := s.dataManager.BatchWrite(c.Request.Context(), convertBatch(batch), time.Now()); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "Failed to batch write data: " + err.Error(),
		})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"status": "ok",
		"count":  len(batch),
	})
}

// handleQuery 处理查询请求
func (s *RESTServer) handleQuery(c *gin.Context) {
	var req QueryRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"error": "Invalid request: " + err.Error(),
		})
		return
	}

	// 创建数据点ID
	id := model.DataPointID{
		DeviceID:   req.DeviceID,
		MetricCode: req.MetricCode,
		Labels:     req.Labels,
	}

	// 创建查询
	queryType := model.QueryType(req.QueryType)
	query := model.NewQuery(queryType, req.Params)

	// 执行查询
	result, err := s.dataManager.ExecuteQuery(c.Request.Context(), id, query)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "Failed to execute query: " + err.Error(),
		})
		return
	}

	// 根据查询类型返回结果
	switch queryType {
	case model.QueryTypeLatest:
		value, ok := result.AsLatest()
		if !ok {
			c.JSON(http.StatusNotFound, gin.H{
				"error": "Data point not found",
			})
			return
		}
		c.JSON(http.StatusOK, gin.H{
			"timestamp": value.Timestamp,
			"value":     value.Value,
		})

	case model.QueryTypeAll:
		values, ok := result.AsAll()
		if !ok {
			c.JSON(http.StatusNotFound, gin.H{
				"error": "Data point not found",
			})
			return
		}
		c.JSON(http.StatusOK, gin.H{
			"values": values,
		})

	case model.QueryTypeDuration:
		duration, ok := result.AsDuration()
		if !ok {
			c.JSON(http.StatusNotFound, gin.H{
				"error": "Data point not found",
			})
			return
		}
		c.JSON(http.StatusOK, gin.H{
			"duration": duration,
		})

	default:
		c.JSON(http.StatusBadRequest, gin.H{
			"error": "Unsupported query type",
		})
	}
}

// Start 启动REST API服务
func (s *RESTServer) Start(addr string) error {
	s.server = &http.Server{
		Addr:    addr,
		Handler: s.router,
	}

	return s.server.ListenAndServe()
}

// Stop 停止REST API服务
func (s *RESTServer) Stop(ctx context.Context) error {
	return s.server.Shutdown(ctx)
}

// convertBatch 将内部批处理格式转换为DataManager.BatchWrite所需的格式
func convertBatch(batch []struct {
	ID    model.DataPointID
	Value model.DataValue
}) []struct {
	ID    model.DataPointID
	Value interface{}
} {
	result := make([]struct {
		ID    model.DataPointID
		Value interface{}
	}, len(batch))

	for i, item := range batch {
		result[i].ID = item.ID
		result[i].Value = item.Value.Value
	}

	return result
}