package agentd

import (
	"context"
	"strings"
	"sync"
	"time"

	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
	"github.com/wangjia/pangolin/server/internal/agentd/v2rayapi"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
)

// V2RayUsageSource reads PER-USER traffic from sing-box's v2ray_api StatsService
// (loopback gRPC). Each provisioned dp_uuid has its own counters
//
//	user>>>{dp_uuid}>>>traffic>>>uplink   (bytes the user sent)
//	user>>>{dp_uuid}>>>traffic>>>downlink (bytes the user received)
//
// QueryStats with reset=true drains+zeroes the counters, so each call yields the
// exact per-user delta for the elapsed window — accurate for any number of
// concurrent users (unlike the old clash node-total split, see ClashUsageSource).
//
// uplink/downlink are already user-perspective (no swap needed). Only opaque
// dp_uuids and byte counts are read — never identities/addresses/DNS.
type V2RayUsageSource struct {
	addr string

	mu     sync.Mutex
	conn   *grpc.ClientConn
	client *v2rayapi.StatsClient
}

func newV2RayUsageSource() *V2RayUsageSource {
	return &V2RayUsageSource{addr: v2rayAPIAddr}
}

// dial lazily establishes (and caches) the loopback gRPC connection.
func (s *V2RayUsageSource) ensureClient() (*v2rayapi.StatsClient, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.client != nil {
		return s.client, nil
	}
	conn, err := grpc.NewClient(s.addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return nil, err
	}
	s.conn = conn
	s.client = v2rayapi.NewStatsClient(conn)
	return s.client, nil
}

// Collect drains per-user counters since the previous call.
func (s *V2RayUsageSource) Collect() []*agentv1.UsageEntry {
	client, err := s.ensureClient()
	if err != nil {
		return nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	resp, err := client.QueryStats(ctx, &v2rayapi.QueryStatsRequest{
		Pattern: "user>>>",
		Reset_:  true,
	})
	if err != nil || resp == nil {
		return nil
	}

	// Aggregate uplink/downlink per dp_uuid.
	type acc struct{ up, down int64 }
	byUUID := make(map[string]*acc)
	for _, st := range resp.GetStat() {
		uuid, dir, ok := parseUserStat(st.GetName())
		if !ok {
			continue
		}
		a := byUUID[uuid]
		if a == nil {
			a = &acc{}
			byUUID[uuid] = a
		}
		switch dir {
		case "uplink":
			a.up += st.GetValue()
		case "downlink":
			a.down += st.GetValue()
		}
	}

	entries := make([]*agentv1.UsageEntry, 0, len(byUUID))
	for uuid, a := range byUUID {
		if a.up == 0 && a.down == 0 {
			continue
		}
		entries = append(entries, &agentv1.UsageEntry{
			DpUUID:         uuid,
			BytesUp:        a.up,
			BytesDown:      a.down,
			SessionMinutes: 1, // 本窗口有流量 → 计 1 分钟在线(近似)
		})
	}
	return entries
}

// parseUserStat splits "user>>>{uuid}>>>traffic>>>{uplink|downlink}".
func parseUserStat(name string) (uuid, dir string, ok bool) {
	parts := strings.Split(name, ">>>")
	if len(parts) != 4 || parts[0] != "user" || parts[2] != "traffic" {
		return "", "", false
	}
	return parts[1], parts[3], true
}