package xray

import (
	"context"
	"net"
	"testing"

	"google.golang.org/grpc"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protodesc"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/types/descriptorpb"
	"google.golang.org/protobuf/types/dynamicpb"
)

// This file proves the wire-level stats client against a STOCK gRPC server:
// the server side runs grpc's default protobuf codec with dynamicpb messages
// built from the xray.app.stats.command descriptor (rebuilt here from the
// .proto field numbers), so what interops is the real wire contract — field
// numbers, message shapes, the "proto" content subtype — not just this
// package's own codec talking to itself.

func fieldDesc(name string, num int32, typ descriptorpb.FieldDescriptorProto_Type, label descriptorpb.FieldDescriptorProto_Label, typeName string) *descriptorpb.FieldDescriptorProto {
	f := &descriptorpb.FieldDescriptorProto{
		Name:   proto.String(name),
		Number: proto.Int32(num),
		Type:   typ.Enum(),
		Label:  label.Enum(),
	}
	if typeName != "" {
		f.TypeName = proto.String(typeName)
	}
	return f
}

func statsFileDescriptor(t *testing.T) protoreflect.FileDescriptor {
	t.Helper()
	const (
		strT  = descriptorpb.FieldDescriptorProto_TYPE_STRING
		i64T  = descriptorpb.FieldDescriptorProto_TYPE_INT64
		boolT = descriptorpb.FieldDescriptorProto_TYPE_BOOL
		msgT  = descriptorpb.FieldDescriptorProto_TYPE_MESSAGE
		opt   = descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL
		rep   = descriptorpb.FieldDescriptorProto_LABEL_REPEATED
		pkg   = ".xray.app.stats.command."
	)
	msg := func(name string, fields ...*descriptorpb.FieldDescriptorProto) *descriptorpb.DescriptorProto {
		return &descriptorpb.DescriptorProto{Name: proto.String(name), Field: fields}
	}
	fdp := &descriptorpb.FileDescriptorProto{
		Name:    proto.String("stats.proto"),
		Package: proto.String("xray.app.stats.command"),
		Syntax:  proto.String("proto3"),
		MessageType: []*descriptorpb.DescriptorProto{
			msg("Stat",
				fieldDesc("name", 1, strT, opt, ""),
				fieldDesc("value", 2, i64T, opt, "")),
			msg("QueryStatsRequest",
				fieldDesc("pattern", 1, strT, opt, ""),
				fieldDesc("reset", 2, boolT, opt, "")),
			msg("QueryStatsResponse",
				fieldDesc("stat", 1, msgT, rep, pkg+"Stat")),
			msg("OnlineIPEntry",
				fieldDesc("ip", 1, strT, opt, ""),
				fieldDesc("last_seen", 2, i64T, opt, "")),
			msg("TrafficUserStat",
				fieldDesc("uplink", 1, i64T, opt, ""),
				fieldDesc("downlink", 2, i64T, opt, "")),
			msg("UserStat",
				fieldDesc("email", 1, strT, opt, ""),
				fieldDesc("ips", 2, msgT, rep, pkg+"OnlineIPEntry"),
				fieldDesc("traffic", 3, msgT, opt, pkg+"TrafficUserStat")),
			msg("GetUsersStatsRequest",
				fieldDesc("include_traffic", 1, boolT, opt, ""),
				fieldDesc("reset", 2, boolT, opt, "")),
			msg("GetUsersStatsResponse",
				fieldDesc("users", 1, msgT, rep, pkg+"UserStat")),
		},
		Service: []*descriptorpb.ServiceDescriptorProto{
			{
				Name: proto.String("StatsService"),
				Method: []*descriptorpb.MethodDescriptorProto{
					{Name: proto.String("QueryStats"), InputType: proto.String(pkg + "QueryStatsRequest"), OutputType: proto.String(pkg + "QueryStatsResponse")},
					{Name: proto.String("GetUsersStats"), InputType: proto.String(pkg + "GetUsersStatsRequest"), OutputType: proto.String(pkg + "GetUsersStatsResponse")},
				},
			},
		},
	}
	fd, err := protodesc.NewFile(fdp, nil)
	if err != nil {
		t.Fatalf("build stats descriptor: %v", err)
	}
	return fd
}

// stockStats is a StatsService served by grpc's stock codec from dynamicpb
// messages, replaying queued responses per method.
type stockStats struct {
	fd       protoreflect.FileDescriptor
	querySeq []*dynamicpb.Message
	usersSeq []*dynamicpb.Message
	qCalls   int
	uCalls   int
}

func (s *stockStats) newMsg(name protoreflect.Name) *dynamicpb.Message {
	return dynamicpb.NewMessage(s.fd.Messages().ByName(name))
}

func (s *stockStats) method(name string, reqName protoreflect.Name, pick func() *dynamicpb.Message) grpc.MethodDesc {
	return grpc.MethodDesc{
		MethodName: name,
		Handler: func(_ any, _ context.Context, dec func(any) error, _ grpc.UnaryServerInterceptor) (any, error) {
			req := s.newMsg(reqName)
			if err := dec(req); err != nil {
				return nil, err
			}
			return pick(), nil
		},
	}
}

func startStockStats(t *testing.T, st *stockStats) *XrayAPI {
	t.Helper()
	st.fd = statsFileDescriptor(t)
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	srv := grpc.NewServer() // stock codec: the real interop surface
	srv.RegisterService(&grpc.ServiceDesc{
		ServiceName: "xray.app.stats.command.StatsService",
		HandlerType: (*statsServiceIface)(nil),
		Methods: []grpc.MethodDesc{
			st.method("QueryStats", "QueryStatsRequest", func() *dynamicpb.Message {
				i := st.qCalls
				st.qCalls++
				if i >= len(st.querySeq) {
					i = len(st.querySeq) - 1
				}
				return st.querySeq[i]
			}),
			st.method("GetUsersStats", "GetUsersStatsRequest", func() *dynamicpb.Message {
				i := st.uCalls
				st.uCalls++
				if i >= len(st.usersSeq) {
					i = len(st.usersSeq) - 1
				}
				return st.usersSeq[i]
			}),
		},
	}, st)
	go func() { _ = srv.Serve(lis) }()
	t.Cleanup(srv.Stop)

	api := &XrayAPI{}
	if err := api.Init(lis.Addr().(*net.TCPAddr).Port); err != nil {
		t.Fatalf("api init: %v", err)
	}
	t.Cleanup(api.Close)
	return api
}

func (s *stockStats) statMsg(name string, value int64) *dynamicpb.Message {
	m := s.newMsg("Stat")
	d := m.Descriptor()
	m.Set(d.Fields().ByName("name"), protoreflect.ValueOfString(name))
	m.Set(d.Fields().ByName("value"), protoreflect.ValueOfInt64(value))
	return m
}

func (s *stockStats) queryResp(stats ...*dynamicpb.Message) *dynamicpb.Message {
	m := s.newMsg("QueryStatsResponse")
	list := m.Mutable(m.Descriptor().Fields().ByName("stat")).List()
	for _, st := range stats {
		list.Append(protoreflect.ValueOfMessage(st.ProtoReflect()))
	}
	return m
}

// TestStatsAgainstStockProtoServer drives GetTraffic and GetOnlineUsers
// against the stock-codec server: the same deltas and user shapes the typed
// client produced must come out of the wire-level one.
func TestStatsAgainstStockProtoServer(t *testing.T) {
	st := &stockStats{}
	st.fd = statsFileDescriptor(t)
	st.querySeq = []*dynamicpb.Message{
		st.queryResp(st.statMsg("user>>>alice>>>traffic>>>uplink", 100)),
		st.queryResp(
			st.statMsg("user>>>alice>>>traffic>>>uplink", 175),
			st.statMsg("inbound>>>in-443>>>traffic>>>downlink", 2048),
		),
	}

	// One online user with two IPs and a traffic block (field 3 — present in
	// the message, ignored by the panel's OnlineUser shape).
	user := st.newMsg("UserStat")
	ud := user.Descriptor()
	user.Set(ud.Fields().ByName("email"), protoreflect.ValueOfString("alice"))
	ipList := user.Mutable(ud.Fields().ByName("ips")).List()
	for _, e := range []struct {
		ip   string
		seen int64
	}{{"203.0.113.7", 111}, {"2001:db8::5", 222}} {
		entry := st.newMsg("OnlineIPEntry")
		ed := entry.Descriptor()
		entry.Set(ed.Fields().ByName("ip"), protoreflect.ValueOfString(e.ip))
		entry.Set(ed.Fields().ByName("last_seen"), protoreflect.ValueOfInt64(e.seen))
		ipList.Append(protoreflect.ValueOfMessage(entry.ProtoReflect()))
	}
	traffic := st.newMsg("TrafficUserStat")
	td := traffic.Descriptor()
	traffic.Set(td.Fields().ByName("uplink"), protoreflect.ValueOfInt64(10))
	traffic.Set(td.Fields().ByName("downlink"), protoreflect.ValueOfInt64(20))
	user.Set(ud.Fields().ByName("traffic"), protoreflect.ValueOfMessage(traffic.ProtoReflect()))
	usersResp := st.newMsg("GetUsersStatsResponse")
	uList := usersResp.Mutable(usersResp.Descriptor().Fields().ByName("users")).List()
	uList.Append(protoreflect.ValueOfMessage(user.ProtoReflect()))
	st.usersSeq = []*dynamicpb.Message{usersResp}

	api := startStockStats(t, st)

	if _, _, err := api.GetTraffic(); err != nil {
		t.Fatalf("GetTraffic (baseline): %v", err)
	}
	tags, clients, err := api.GetTraffic()
	if err != nil {
		t.Fatalf("GetTraffic: %v", err)
	}
	if len(clients) != 1 || clients[0].Email != "alice" || clients[0].Up != 75 {
		t.Fatalf("clients = %+v, want alice up=75", clients)
	}
	if len(tags) != 1 || tags[0].Tag != "in-443" || tags[0].Down != 2048 {
		t.Fatalf("tags = %+v, want in-443 down=2048", tags)
	}

	users, err := api.GetOnlineUsers()
	if err != nil {
		t.Fatalf("GetOnlineUsers: %v", err)
	}
	if len(users) != 1 || users[0].Email != "alice" || len(users[0].IPs) != 2 {
		t.Fatalf("users = %+v, want alice with 2 IPs", users)
	}
	if users[0].IPs[0].IP != "203.0.113.7" || users[0].IPs[0].LastSeen != 111 ||
		users[0].IPs[1].IP != "2001:db8::5" || users[0].IPs[1].LastSeen != 222 {
		t.Fatalf("ips = %+v", users[0].IPs)
	}
}