package throttler

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

	"github.com/pingcap/log"
	"github.com/tidbcloud/ffs/pkg/v1/proto"
	uclient "github.com/tidbcloud/ffs/universe/client"
	zclient "github.com/tidbcloud/ffs/zone/client"
	"golang.org/x/time/rate"
)

type throttleMeta struct {
	*proto.ThrottleReq
	mu struct {
		sync.RWMutex
		ReadBytesLimiter  *rate.Limiter
		ReadIoLimiter     *rate.Limiter
		WriteBytesLimiter *rate.Limiter
		WriteIoLimiter    *rate.Limiter
	}
}

type ThrottlerServer struct {
	proto.UnimplementedThrottlingServiceServer
	ctx       context.Context
	cancel    context.CancelFunc
	uclient   *uclient.Client
	zclient   *zclient.Client
	throttles struct {
		sync.RWMutex
		metas map[string]*throttleMeta
	}
}

func NewThrottlerServer(universe, zone, address string) *ThrottlerServer {
	callCtx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()

	uclient := uclient.NewClient(strings.Split(universe, ",")...)
	if err := uclient.Start(); err != nil {
		log.S().Fatalf("failed to start universe client: %v", err)
	}
	var z *proto.Zone
	if resp, err := uclient.Client().GetZone(callCtx, &proto.GetZoneReq{Name: zone}); err != nil {
		log.S().Fatalf("failed to get zone: %v", err)
	} else if len(resp.Endpoints) == 0 {
		log.S().Fatalf("failed to get zone %v endpoint", resp.Name)
	} else {
		log.S().Infof("get zone name: %s endpoints: %v", resp.Name, resp.Endpoints)
		z = resp
	}

	zclient := zclient.NewClient(z.Endpoints...)
	if err := zclient.Start(); err != nil {
		log.S().Fatalf("failed to start zone client: %v", err)
	}
	if _, err := zclient.Client().ThrottlerHeartbeat(callCtx, &proto.ThrottlerHeartbeatReq{Throttler: &proto.Throttler{Address: address}}); err != nil {
		log.S().Fatalf("failed to heartbeat: %v", err)
	}
	svr := &ThrottlerServer{uclient: uclient, zclient: zclient}
	svr.throttles.metas = make(map[string]*throttleMeta)
	svr.ctx, svr.cancel = context.WithCancel(context.Background())
	go func() {
		for {
			stream, err := svr.uclient.Client().WatchZones(svr.ctx, &proto.WatchZonesReq{Pattern: z.Name})
			if err != nil {
				log.S().Errorf("failed to watch zones: %v", err)
			} else {
				for {
					zonesList, err := stream.Recv()
					if err != nil {
						log.S().Warn(err)
						break
					}
					for _, zone := range zonesList.Zones {
						if zone.Name == z.Name {
							svr.zclient.Update(zone.Endpoints...)
							log.S().Warnf("update zone %s endpoints %v", z.Name, zone.Endpoints)
							break
						}
					}
				}
			}
			time.Sleep(time.Second)
		}
	}()
	go func() {
		ticker := time.NewTicker(time.Second)
		for {
			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
			if _, err := svr.zclient.Client().ThrottlerHeartbeat(ctx, &proto.ThrottlerHeartbeatReq{Throttler: &proto.Throttler{Address: address}}); err != nil {
				log.S().Errorf("failed to heartbeat: %v", err)
			}
			cancel()
			select {
			case <-ticker.C:
			case <-svr.ctx.Done():
				ticker.Stop()
				return
			}
		}
	}()
	return svr
}

func (s *ThrottlerServer) Stop() {
	s.uclient.Stop()
	s.zclient.Stop()
}

func (s *ThrottlerServer) Throttle(ctx context.Context, req *proto.ThrottleReq) (*proto.ThrottleResp, error) {
	s.throttles.RLock()
	m, ok := s.throttles.metas[req.GetTenant()]
	s.throttles.RUnlock()
	if ok {
		m.mu.RLock()
		if req.ReadBytes > 0 {
			m.mu.ReadBytesLimiter.WaitN(ctx, int(req.ReadBytes))
		}
		if req.ReadIo > 0 {
			m.mu.ReadIoLimiter.WaitN(ctx, int(req.ReadIo))
		}
		if req.WriteBytes > 0 {
			m.mu.WriteBytesLimiter.WaitN(ctx, int(req.WriteBytes))
		}
		if req.WriteIo > 0 {
			m.mu.WriteIoLimiter.WaitN(ctx, int(req.WriteIo))
		}
		m.mu.RUnlock()
		return &proto.ThrottleResp{
			ReadBytes:  req.ReadBytes,
			ReadIo:     req.ReadIo,
			WriteBytes: req.WriteBytes,
			WriteIo:    req.WriteIo,
		}, nil
	}
	return nil, errors.New("tenant throttle not found")
}

func (s *ThrottlerServer) Handover(ctx context.Context, req *proto.HandoverReq) (*proto.EmptyMessage, error) {
	tenants := req.Tenants
	for _, t := range tenants {
		s.throttles.Lock()
		if m, ok := s.throttles.metas[t.Tenant]; ok {
			old := m.ThrottleReq
			m.ThrottleReq = t
			s.throttles.Unlock()
			m.mu.Lock()
			if old.ReadBytes != t.ReadBytes {
				m.mu.ReadBytesLimiter = rate.NewLimiter(rate.Limit(t.ReadBytes), int(t.ReadBytes))
			}
			if old.ReadIo != t.ReadIo {
				m.mu.ReadIoLimiter = rate.NewLimiter(rate.Limit(t.ReadIo), int(t.ReadIo))
			}
			if old.WriteBytes != t.WriteBytes {
				m.mu.WriteBytesLimiter = rate.NewLimiter(rate.Limit(t.WriteBytes), int(t.WriteBytes))
			}
			if old.WriteIo != t.WriteIo {
				m.mu.WriteIoLimiter = rate.NewLimiter(rate.Limit(t.WriteIo), int(t.WriteIo))
			}
			m.mu.Unlock()
		} else {
			m := &throttleMeta{ThrottleReq: t}
			m.mu.ReadBytesLimiter = rate.NewLimiter(rate.Limit(t.ReadBytes), int(t.ReadBytes))
			m.mu.ReadIoLimiter = rate.NewLimiter(rate.Limit(t.ReadIo), int(t.ReadIo))
			m.mu.WriteBytesLimiter = rate.NewLimiter(rate.Limit(t.WriteBytes), int(t.WriteBytes))
			m.mu.WriteIoLimiter = rate.NewLimiter(rate.Limit(t.WriteIo), int(t.WriteIo))
			s.throttles.metas[t.Tenant] = m
			s.throttles.Unlock()
		}
	}
	return nil, nil
}