package app

import (
	"io"
	"net"

	"github.com/golang/protobuf/proto"
	kvpb "github.com/pingcap/kvproto/pkg/resource_usage_agent"
	"github.com/pingcap/log"
	"github.com/pingcap/tipb/go-tipb"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

var (
	_ tipb.TopSQLAgentServer        = &TopSQLAgentServer{}
	_ kvpb.ResourceUsageAgentServer = &TopSQLAgentServer{}
)

// TopSQLAgentServer is the main struct which controls the Top SQL metrics collecting and reporting.
//
// In general, there are several parts in the agent which are working together:
// - gRPC server receives raw data from tidb-server/tikv-server
// - send buffer writes gRPC data to an in-memory buffer, which is a sliding-window-like
//   in-memory view of the data. The send buffer has an upper limit, which can be seemed as
//   the max window size of the view.
// - sender reads data from send buffer, and asynchronously sends it to the remote store.
type TopSQLAgentServer struct {
	sender  *Sender
	sendBuf *SendBuffer
}

func NewAgentServer(store Store) *TopSQLAgentServer {
	// TODO(dragonly): make this configurable
	sendBuf := NewSendBuffer(SendBufferConfig{
		MaxCPUTimeRecordCount: 1000,
		MaxSQLMetaCount:       100,
		MaxPlanMetaCount:      100,
	})
	sender := NewSender(sendBuf, store)
	return &TopSQLAgentServer{
		sender:  sender,
		sendBuf: sendBuf,
	}
}

func (s *TopSQLAgentServer) Start() {
	s.sender.Start()
}

// ReportPlanMeta implements RPCs with tidb-server
func (s *TopSQLAgentServer) ReportPlanMeta(stream tipb.TopSQLAgent_ReportPlanMetaServer) error {
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		log.Debug("ReportPlanMeta[tidb]", zap.Any("meta", req))
		s.receivePlanMeta(req)
	}
	resp := &tipb.EmptyResponse{}
	return stream.SendAndClose(resp)
}

// ReportSQLMeta implements RPCs with tidb-server
func (s *TopSQLAgentServer) ReportSQLMeta(stream tipb.TopSQLAgent_ReportSQLMetaServer) error {
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		log.Debug("ReportSQLMeta[tidb]", zap.Any("meta", req))
		s.receiveSQLMeta(req)
	}
	resp := &tipb.EmptyResponse{}
	return stream.SendAndClose(resp)
}

// ReportCPUTimeRecords implements RPCs with tidb-server
func (s *TopSQLAgentServer) ReportCPUTimeRecords(stream tipb.TopSQLAgent_ReportCPUTimeRecordsServer) error {
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		log.Debug("ReportCPUTimeRecords[tidb]",
			zap.Int("timestamp list size", len(req.RecordListTimestampSec)),
			zap.Uint32s("cpu time list", req.RecordListCpuTimeMs),
			zap.ByteString("sql digest", req.SqlDigest),
			zap.ByteString("plan digest", req.PlanDigest),
		)
		s.receiveCPUTimeRecords(req)
	}
	resp := &tipb.EmptyResponse{}
	return stream.SendAndClose(resp)
}

// ReportCPUTime implements RPCs with tikv-server
func (s *TopSQLAgentServer) ReportCPUTime(stream kvpb.ResourceUsageAgent_ReportCPUTimeServer) error {
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		s.report(req.ResourceGroupTag, req.RecordListCpuTimeMs, req.GetRecordListTimestampSec())
	}
	resp := &kvpb.EmptyResponse{}
	stream.SendAndClose(resp)
	return nil
}

// Report implements RPCs with tikv-server.
// Currently, this is only used to report CPU time.
func (s *TopSQLAgentServer) Report(stream kvpb.ResourceUsageAgent_ReportServer) error {
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		s.report(req.ResourceGroupTag, req.RecordListCpuTimeMs, req.GetRecordListTimestampSec())
	}
	resp := &kvpb.EmptyResponse{}
	stream.SendAndClose(resp)
	return nil
}

func (s *TopSQLAgentServer) report(tag []byte, cpuTime []uint32, timestamp []uint64) {
	resourceGroupTag := &tipb.ResourceGroupTag{}
	if err := proto.Unmarshal(tag, resourceGroupTag); err != nil {
		log.Error("failed to Unmarshal resourceGroupTag", zap.Error(err))
	}
	record := &tipb.CPUTimeRecord{
		SqlDigest:              resourceGroupTag.SqlDigest,
		PlanDigest:             resourceGroupTag.PlanDigest,
		RecordListTimestampSec: timestamp,
		RecordListCpuTimeMs:    cpuTime,
	}
	log.Debug("ReportCPUTime[tikv]",
		zap.Int("timestamp list size", len(timestamp)),
		zap.Uint32s("cpu time list", cpuTime),
		zap.ByteString("sql digest", resourceGroupTag.SqlDigest),
		zap.ByteString("plan digest", resourceGroupTag.PlanDigest),
	)
	s.receiveCPUTimeRecords(record)
}

func (s *TopSQLAgentServer) receiveCPUTimeRecords(record *tipb.CPUTimeRecord) {
	if len(record.RecordListTimestampSec) > 0 {
		s.sendBuf.WriteOneCPUTimeRecordOrDrop(record)
	}
}

func (s *TopSQLAgentServer) receiveSQLMeta(meta *tipb.SQLMeta) {
	s.sendBuf.WriteOneSQLMetaOrDrop(meta)
}

func (s *TopSQLAgentServer) receivePlanMeta(meta *tipb.PlanMeta) {
	s.sendBuf.WriteOnePlanMetaOrDrop(meta)
}

type Server struct {
	address     string
	store       Store
	agentServer *TopSQLAgentServer
	grpcServer  *grpc.Server
	tcpListener net.Listener
}

func NewServer(address string, store Store) *Server {
	s := &Server{
		address: address,
		store:   store,
	}
	return s
}

func (s *Server) Start() {
	tcpListener, err := net.Listen("tcp", s.address)
	if err != nil {
		log.Error("failed to listen on tcp address", zap.String("address", s.address), zap.Error(err))
		return
	}
	s.tcpListener = tcpListener
	s.grpcServer = grpc.NewServer()
	s.agentServer = NewAgentServer(s.store)
	go s.agentServer.Start()
	tipb.RegisterTopSQLAgentServer(s.grpcServer, s.agentServer)
	kvpb.RegisterResourceUsageAgentServer(s.grpcServer, s.agentServer)

	if err := s.grpcServer.Serve(s.tcpListener); err != nil {
		log.Error("failed to start gRPC server", zap.Error(err))
	}
}