package grpc_conn

import (
	"context"
	"fmt"
	"google.golang.org/grpc/stats"
)

// 1. 实现 stats.Handler 接口
type StatsHandler struct{}

func (h *StatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
	// 给 RPC 调用打标签（例如记录方法名）
	return context.WithValue(ctx, "rpc_method", info.FullMethodName)
}

func (h *StatsHandler) HandleRPC(ctx context.Context, s stats.RPCStats) {
	// 处理 RPC 统计信息
	switch t := s.(type) {
	case *stats.Begin:
		fmt.Printf("RPC started: %s\n", ctx.Value("rpc_method"))
	case *stats.End:
		fmt.Printf("RPC finished: %s (duration: %v)\n",
			ctx.Value("rpc_method"), t.EndTime.Sub(t.BeginTime))
	case *stats.InPayload:
		fmt.Printf("Received %d bytes\n", t.WireLength)
	case *stats.OutPayload:
		fmt.Printf("Sent %d bytes\n", t.WireLength)
	}
}

func (h *StatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
	// 给连接打标签
	return ctx
}

func (h *StatsHandler) HandleConn(ctx context.Context, s stats.ConnStats) {
	// 处理连接事件
	switch s.(type) {
	case *stats.ConnBegin:
		fmt.Println("Connection established")
	case *stats.ConnEnd:
		fmt.Println("Connection closed")
	}
}