gotidb/pkg/client/quic_client.go

package client

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"time"

	"github.com/quic-go/quic-go"
)

// QUICClient represents a QUIC protocol client
type QUICClient struct {
	conn       quic.Connection
	tlsConfig  *tls.Config
	quicConfig *quic.Config
	address    string
}

// WriteRequest represents a write request to the server
type WriteRequest struct {
	Points  []DataPoint `json:"points"`
	BatchID string      `json:"batch_id"`
}

// DataPoint represents a single data point
type DataPoint struct {
	DeviceID   string            `json:"device_id"`
	MetricCode string            `json:"metric_code"`
	Labels     map[string]string `json:"labels"`
	Value      float64           `json:"value"`
	Timestamp  int64             `json:"timestamp"`
}

// WriteResponse represents the response from the server
type WriteResponse struct {
	Success     bool   `json:"success"`
	Error       string `json:"error,omitempty"`
	BatchID     string `json:"batch_id"`
	PointsCount int    `json:"points_count"`
}

// NewQUICClient creates a new QUIC client
func NewQUICClient(address string) (*QUICClient, error) {
	tlsConfig := &tls.Config{
		InsecureSkipVerify: true, // For development only
		NextProtos:         []string{"gotidb-quic"},
	}

	quicConfig := &quic.Config{
		MaxIdleTimeout:  30 * time.Second,
		KeepAlivePeriod: 10 * time.Second,
	}

	return &QUICClient{
		address:    address,
		tlsConfig:  tlsConfig,
		quicConfig: quicConfig,
	}, nil
}

// Connect establishes a connection to the QUIC server
func (c *QUICClient) Connect(ctx context.Context) error {
	conn, err := quic.DialAddr(ctx, c.address, c.tlsConfig, c.quicConfig)
	if err != nil {
		return fmt.Errorf("failed to connect to QUIC server: %v", err)
	}
	c.conn = conn
	return nil
}

// Close closes the connection to the QUIC server
func (c *QUICClient) Close() error {
	if c.conn != nil {
		return c.conn.CloseWithError(0, "normal closure")
	}
	return nil
}

// Write sends data points to the server
func (c *QUICClient) Write(ctx context.Context, points []DataPoint) (*WriteResponse, error) {
	if c.conn == nil {
		return nil, fmt.Errorf("client not connected")
	}

	// Open a new stream for this write request
	stream, err := c.conn.OpenStreamSync(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to open stream: %v", err)
	}
	defer stream.Close()

	// Create and send the write request
	req := WriteRequest{
		Points:  points,
		BatchID: fmt.Sprintf("%d", time.Now().UnixNano()),
	}

	encoder := json.NewEncoder(stream)
	if err := encoder.Encode(req); err != nil {
		return nil, fmt.Errorf("failed to encode request: %v", err)
	}

	// Read the response
	var resp WriteResponse
	decoder := json.NewDecoder(stream)
	if err := decoder.Decode(&resp); err != nil {
		return nil, fmt.Errorf("failed to decode response: %v", err)
	}

	return &resp, nil
}

// WriteBatch is a convenience method for writing multiple data points
func (c *QUICClient) WriteBatch(ctx context.Context, deviceID, metricCode string, values []float64, timestamps []int64, labels map[string]string) (*WriteResponse, error) {
	if len(values) != len(timestamps) {
		return nil, fmt.Errorf("values and timestamps length mismatch")
	}

	points := make([]DataPoint, len(values))
	for i := range values {
		points[i] = DataPoint{
			DeviceID:   deviceID,
			MetricCode: metricCode,
			Labels:     labels,
			Value:      values[i],
			Timestamp:  timestamps[i],
		}
	}

	return c.Write(ctx, points)
}

// Example usage:
/*
func main() {
	client, err := NewQUICClient("localhost:4242")
	if err != nil {
		log.Fatal(err)
	}

	ctx := context.Background()
	if err := client.Connect(ctx); err != nil {
		log.Fatal(err)
	}
	defer client.Close()

	points := []DataPoint{
		{
			DeviceID:   "device1",
			MetricCode: "temperature",
			Labels:     map[string]string{"location": "room1"},
			Value:      25.5,
			Timestamp:  time.Now().UnixNano(),
		},
	}

	resp, err := client.Write(ctx, points)
	if err != nil {
		log.Fatal(err)
	}

	if !resp.Success {
		log.Printf("Write failed: %s", resp.Error)
	} else {
		log.Printf("Successfully wrote %d points", resp.PointsCount)
	}
}
*/