package app

import (
	"context"
	"log"
	"testing"
	"time"

	"github.com/pingcap/tipb/go-tipb"
	"github.com/stretchr/testify/assert"
	"google.golang.org/grpc"
)

func TestSendData(t *testing.T) {
	addr := "localhost:23333"
	store := NewMemStore()
	server := NewServer(addr, store)
	go server.Start()
	time.Sleep(1 * time.Second)
	conn, client := newGrpcClient(context.TODO(), addr)
	defer conn.Close()

	// CPU time record
	var cpuTimeRecordBatch []*tipb.CPUTimeRecord
	count := 1000
	for i := 0; i < count; i++ {
		cpuTimeRecordBatch = append(cpuTimeRecordBatch,
			&tipb.CPUTimeRecord{
				SqlDigest:              []byte("SQLDigest"),
				PlanDigest:             []byte("PlanDigest"),
				RecordListTimestampSec: []uint64{uint64(i)},
				RecordListCpuTimeMs:    []uint32{uint32(i)},
			},
		)
	}
	client.sendCPUTimeRecordBatch(cpuTimeRecordBatch)
	// wait sender to send out the record
	time.Sleep(1 * time.Second)
	cpuTimeRecordList := server.agentServer.sender.store.(*MemStore).cpuTimeRecordList
	assert.Equal(t, count, len(cpuTimeRecordList))
	for i := 0; i < count; i++ {
		assert.Equal(t, *cpuTimeRecordBatch[i], *cpuTimeRecordList[i])
	}

	// SQL meta
	var sqlMetaBatch []*tipb.SQLMeta
	count = 100
	for i := 0; i < count; i++ {
		sqlMetaBatch = append(sqlMetaBatch,
			&tipb.SQLMeta{
				SqlDigest:     []byte("SQLDigest"),
				NormalizedSql: "NormalizedSQL",
			},
		)
	}
	client.sendSQLMetaBatch(sqlMetaBatch)
	time.Sleep(1 * time.Second)
	sqlMetaList := server.agentServer.sender.store.(*MemStore).sqlMetaList
	assert.Equal(t, count, len(sqlMetaList))
	for i := 0; i < count; i++ {
		assert.Equal(t, *sqlMetaBatch[i], *sqlMetaList[i])
	}

	// Plan meat
	var planMetaBatch []*tipb.PlanMeta
	count = 100
	for i := 0; i < count; i++ {
		planMetaBatch = append(planMetaBatch,
			&tipb.PlanMeta{
				PlanDigest:     []byte("PlanDigest"),
				NormalizedPlan: "NormalizedPlan",
			},
		)
	}
	client.sendPlanMetaBatch(planMetaBatch)
	time.Sleep(1 * time.Second)
	planMetaList := server.agentServer.sender.store.(*MemStore).planMetaList
	assert.Equal(t, count, len(planMetaList))
	for i := 0; i < count; i++ {
		assert.Equal(t, *planMetaBatch[i], *planMetaList[i])
	}
}

type Client struct {
	cpuTimeRecordStream tipb.TopSQLAgent_ReportCPUTimeRecordsClient
	sqlMetaStream       tipb.TopSQLAgent_ReportSQLMetaClient
	planMetaStream      tipb.TopSQLAgent_ReportPlanMetaClient
}

func NewClient(
	cpuTimeRecordStream tipb.TopSQLAgent_ReportCPUTimeRecordsClient,
	sqlMetaStream tipb.TopSQLAgent_ReportSQLMetaClient,
	planMetaStream tipb.TopSQLAgent_ReportPlanMetaClient,
) *Client {
	return &Client{
		cpuTimeRecordStream: cpuTimeRecordStream,
		sqlMetaStream:       sqlMetaStream,
		planMetaStream:      planMetaStream,
	}
}

// Start starts a goroutine, which sends tidb-server's last minute's data to the gRPC server
func (s *Client) Start() {
	var reqBatch []*tipb.CPUTimeRecord
	for i := 0; i < 10; i++ {
		req := &tipb.CPUTimeRecord{
			SqlDigest:              []byte("SQLDigest"),
			PlanDigest:             []byte("PlanDigest"),
			RecordListTimestampSec: []uint64{uint64(i)},
			RecordListCpuTimeMs:    []uint32{uint32(i * 100)},
		}
		reqBatch = append(reqBatch, req)
	}

	s.sendCPUTimeRecordBatch(reqBatch)
}

func (s *Client) sendCPUTimeRecordBatch(batch []*tipb.CPUTimeRecord) {
	for _, req := range batch {
		if err := s.cpuTimeRecordStream.Send(req); err != nil {
			log.Fatalf("send stream request failed: %v", err)
		}
	}
	resp, err := s.cpuTimeRecordStream.CloseAndRecv()
	if err != nil {
		log.Fatalf("receive stream response failed: %v", err)
	}
	log.Printf("received stream response: %v", resp)
}

func (s *Client) sendSQLMetaBatch(batch []*tipb.SQLMeta) {
	for _, req := range batch {
		if err := s.sqlMetaStream.Send(req); err != nil {
			log.Fatalf("send stream request failed: %v", err)
		}
	}
	resp, err := s.sqlMetaStream.CloseAndRecv()
	if err != nil {
		log.Fatalf("receive stream response failed: %v", err)
	}
	log.Printf("received stream response: %v", resp)
}

func (s *Client) sendPlanMetaBatch(batch []*tipb.PlanMeta) {
	for _, req := range batch {
		if err := s.planMetaStream.Send(req); err != nil {
			log.Fatalf("send stream request failed: %v", err)
		}
	}
	resp, err := s.planMetaStream.CloseAndRecv()
	if err != nil {
		log.Fatalf("receive stream response failed: %v", err)
	}
	log.Printf("received stream response: %v", resp)
}

func newGrpcClient(ctx context.Context, addr string) (*grpc.ClientConn, *Client) {
	dialCtx, cancel := context.WithTimeout(context.TODO(), time.Second)
	defer cancel()
	conn, err := grpc.DialContext(dialCtx, addr, grpc.WithInsecure())
	if err != nil {
		log.Printf("connecting server failed: %v", err)
	}
	grpcClient := tipb.NewTopSQLAgentClient(conn)
	cpuTimeRecordStream, err := grpcClient.ReportCPUTimeRecords(ctx)
	if err != nil {
		log.Printf("open cpu time record stream failed: %v", err)
	}
	sqlMetaStream, err := grpcClient.ReportSQLMeta(ctx)
	if err != nil {
		log.Printf("open cpu time record stream failed: %v", err)
	}
	planMetaStream, err := grpcClient.ReportPlanMeta(ctx)
	if err != nil {
		log.Printf("open cpu time record stream failed: %v", err)
	}

	tidbSender := NewClient(
		cpuTimeRecordStream,
		sqlMetaStream,
		planMetaStream,
	)
	return conn, tidbSender
}

func SendRequest() {
	addr := "localhost:23333"
	conn, sender := newGrpcClient(context.TODO(), addr)
	sender.Start()
	conn.Close()
}