package zone

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

	"github.com/hslam/topk"
	"github.com/pingcap/log"
	"github.com/tidbcloud/ffs/config"
	"github.com/tidbcloud/ffs/pkg/v1/proto"
	uclient "github.com/tidbcloud/ffs/universe/client"
	"github.com/tidbcloud/ffs/util/codec"
	"go.etcd.io/etcd/clientv3"
)

func cloneRealm(r *proto.Realm) *proto.Realm {
	b, _ := r.Marshal()
	clone := &proto.Realm{}
	clone.Unmarshal(b)
	return clone
}

type realmMeta struct {
	*proto.Realm
	value []byte
	timer *time.Timer
	mu    struct {
		sync.RWMutex
		watchers map[*realmWatcher]struct{}
	}
}

func (m *realmMeta) broadcast(r *proto.Realm) {
	m.mu.RLock()
	for w := range m.mu.watchers {
		w.notify(r)
	}
	m.mu.RUnlock()
}

type realmWatcher struct {
	ch chan *proto.Realm
}

func (w *realmWatcher) notify(r *proto.Realm) {
	w.ch <- r
}

func cloneThrottler(t *proto.Throttler) *proto.Throttler {
	b, _ := t.Marshal()
	clone := &proto.Throttler{}
	clone.Unmarshal(b)
	return clone
}

type throttlerMeta struct {
	*proto.Throttler
	value []byte
	timer *time.Timer
	mu    struct {
		sync.RWMutex
		watchers map[*throttlerWatcher]struct{}
	}
}

func (m *throttlerMeta) broadcast(t *proto.Throttler) {
	m.mu.RLock()
	for w := range m.mu.watchers {
		w.notify(t)
	}
	m.mu.RUnlock()
}

type throttlerWatcher struct {
	ch chan *proto.Throttler
}

func (w *throttlerWatcher) notify(r *proto.Throttler) {
	w.ch <- r
}

type ZoneServer struct {
	proto.UnimplementedZoneServiceServer
	name     string
	endpoint string
	eclient  *clientv3.Client
	ctx      context.Context
	cancel   context.CancelFunc
	uclient  *uclient.Client
	realms   struct {
		sync.RWMutex
		metas map[string]*realmMeta
	}
	throttlers struct {
		sync.RWMutex
		metas map[string]*throttlerMeta
	}
}

func NewZoneServer(universe, name, addr string, storeClient *clientv3.Client) *ZoneServer {
	uclient := uclient.NewClient(strings.Split(universe, ",")...)
	if err := uclient.Start(); err != nil {
		log.S().Fatalf("failed to start client: %v", err)
	}
	svr := &ZoneServer{
		name:     name,
		endpoint: addr,
		eclient:  storeClient,
		uclient:  uclient,
	}
	svr.realms.metas = make(map[string]*realmMeta)
	svr.throttlers.metas = make(map[string]*throttlerMeta)
	svr.ctx, svr.cancel = context.WithCancel(context.Background())
	go func() {
		ticker := time.NewTicker(time.Second * 3)
		req := &proto.ZoneHeartbeatReq{Zone: &proto.Zone{Name: name, Endpoints: []string{addr}}}
		for {
			ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
			if _, err := uclient.Client().ZoneHeartbeat(ctx, req); err != nil {
				log.S().Errorf("failed to heartbeat: %v", err)
			}
			cancel()
			select {
			case <-ticker.C:
			case <-svr.ctx.Done():
				ticker.Stop()
				return
			}
		}
	}()
	svr.watchRealms()
	svr.watchThrottlers()
	return svr
}

func (s *ZoneServer) Stop() {
	s.uclient.Stop()
	s.cancel()
}

func (s *ZoneServer) watchRealms() {
	resp, err := s.eclient.Get(context.TODO(), config.RealmMetaPrefix, clientv3.WithPrefix())
	if err != nil {
		log.S().Fatalf("failed to get realms: %v", err)
	}
	for _, kv := range resp.Kvs {
		s.getOrCreateRealmMeta(kv.Value)
	}
	rev := resp.Header.Revision + 1
	go func() {
		watcher := clientv3.NewWatcher(s.eclient)
		watchChan := watcher.Watch(s.ctx, config.RealmMetaPrefix, clientv3.WithPrefix(), clientv3.WithRev(rev))
		for watchResponse := range watchChan {
			for _, event := range watchResponse.Events {
				switch event.Type {
				case clientv3.EventTypePut:
					s.getOrCreateRealmMeta(event.Kv.Value)
				case clientv3.EventTypeDelete:
					s.deleteRealmMeta(string(event.Kv.Key)[len(config.RealmMetaPrefix):])
				}
			}
		}
	}()
}

func (s *ZoneServer) getOrCreateRealmMeta(value []byte) *realmMeta {
	var t = &proto.Realm{}
	codec.Unmarshal(value, t)
	s.realms.Lock()
	m, ok := s.realms.metas[t.Name]
	if ok {
		changed := m.Address != t.Address
		m.Realm = t
		m.value = value
		clone := cloneRealm(m.Realm)
		s.realms.Unlock()
		if changed {
			log.S().Infof("update realm meta %v", m.Realm)
			m.broadcast(clone)
		}
	} else {
		m = &realmMeta{Realm: t, value: value}
		m.mu.watchers = make(map[*realmWatcher]struct{})
		s.realms.metas[t.Name] = m
		s.realms.Unlock()
		log.S().Infof("create realm meta %v", m.Realm)
	}
	return m
}

func (s *ZoneServer) deleteRealmMeta(n string) {
	log.S().Infof("delete realm meta %s", n)
	s.realms.Lock()
	if m, ok := s.realms.metas[n]; ok {
		m.Realm = &proto.Realm{}
		delete(s.realms.metas, n)
		clone := cloneRealm(m.Realm)
		s.realms.Unlock()
		m.broadcast(clone)
	} else {
		s.realms.Unlock()
	}
}

func (s *ZoneServer) RealmHeartbeat(ctx context.Context, req *proto.RealmHeartbeatReq) (*proto.RealmHeartbeatResp, error) {
	r := req.Realm
	if r == nil {
		return nil, errors.New("empty realm")
	}
	value, err := codec.Marshal(r)
	if err != nil {
		return nil, err
	}
	if _, err := s.eclient.Put(ctx, config.RealmMetaPrefix+r.Name, string(value)); err != nil {
		log.S().Errorf("realm heartbeat: name %s address %s err", req.Realm.Name, req.Realm.Address, err)
		return nil, err
	}
	numKey := config.RealmNumClientsPrefix + r.Name
	if resp, err := s.eclient.Get(ctx, numKey); err != nil {
		log.S().Error(err)
		return nil, err
	} else if len(resp.Kvs) == 0 {
		txn := s.eclient.Txn(context.TODO())
		txn.If(clientv3.Compare(clientv3.Version(numKey), "=", 0)).
			Then(clientv3.OpPut(numKey, strconv.Itoa(0))).
			Else()
		_, err := txn.Commit()
		if err != nil {
			log.S().Error(err)
			return nil, err
		}
	}
	return &proto.RealmHeartbeatResp{}, nil
}

type realmNumClients struct {
	Realm       string
	NumClients  int
	ModRevision int64
}

type list []*realmNumClients

func (l list) Len() int { return len(l) }
func (l list) Less(i, j int) bool {
	return l[i].NumClients > l[j].NumClients
}
func (l list) Swap(i, j int) { l[i], l[j] = l[j], l[i] }

func (s *ZoneServer) assign(policy *proto.StoragePolicy, merged bool) ([]*realmNumClients, error) {
	resp, err := s.eclient.Get(context.TODO(), config.RealmNumClientsPrefix, clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	var l list
	for _, kv := range resp.Kvs {
		var t = &realmNumClients{}
		t.Realm = string(kv.Key)[len(config.RealmNumClientsPrefix):]
		t.NumClients, _ = strconv.Atoi(string(kv.Value))
		t.ModRevision = kv.ModRevision
		l = append(l, t)
	}

	if merged && len(l) > 1 {
		topk.Top(l, 2, false)
		return l[:2], nil
	} else if len(l) > 0 {
		topk.Top(l, 1, false)
		return l[:1], nil
	}
	return nil, errors.New("no realms available")
}

func (s *ZoneServer) Locate(ctx context.Context, req *proto.LocateReq) (*proto.MountPoint, error) {
	log.S().Infof("locate tenant %s mount point %s", req.Tenant, req.MountPoint)
	for {
		mnt, err := s.locate(ctx, req)
		if mnt == nil && err == nil {
			continue
		}
		return mnt, err
	}
}

func (s *ZoneServer) locate(ctx context.Context, req *proto.LocateReq) (*proto.MountPoint, error) {
	key := config.MountPrefix + req.Tenant + "/" + req.MountPoint
	strings.ReplaceAll(key, "//", "/")
	resp, err := s.eclient.Get(ctx, key)
	if err != nil {
		log.S().Error(err)
		return nil, err
	}
	if len(resp.Kvs) > 0 && len(resp.Kvs[0].Value) > 0 {
		names := strings.Split(string(resp.Kvs[0].Value), ",")
		realms := make([]*proto.Realm, len(names))
		for i, name := range names {
			r, err := s.GetRealm(ctx, &proto.GetRealmReq{Name: name})
			if err != nil {
				log.S().Errorf("Get Realm %s err %s", name, err)
				return nil, err
			}
			realms[i] = r
		}
		log.S().Infof("locate tenant %s mount point %s successfully %v", req.Tenant, req.MountPoint, realms)
		return &proto.MountPoint{Name: req.MountPoint, Realms: realms}, nil
	} else if !req.Create {
		return &proto.MountPoint{Name: req.MountPoint}, nil
	}
	realmNumClients, err := s.assign(req.Policy, req.Merged)
	if err != nil {
		log.S().Error(err)
		return nil, err
	}
	var names = make([]string, len(realmNumClients))
	for i, r := range realmNumClients {
		names[i] = r.Realm
	}
	sort.Strings(names)
	value := strings.Join(names, ",")
	cmps := []clientv3.Cmp{
		clientv3.Compare(clientv3.Version(key), "=", 0),
	}
	ops := []clientv3.Op{
		clientv3.OpPut(key, value),
	}
	for _, r := range realmNumClients {
		numKey := config.RealmNumClientsPrefix + r.Realm
		cmp := clientv3.Compare(clientv3.ModRevision(numKey), "=", r.ModRevision)
		cmps = append(cmps, cmp)
		r.NumClients++
		op := clientv3.OpPut(numKey, strconv.Itoa(r.NumClients))
		ops = append(ops, op)
	}
	txn := s.eclient.Txn(context.TODO())
	txn.If(cmps...).
		Then(ops...).
		Else(clientv3.OpGet(key))
	txnResp, err := txn.Commit()
	if err != nil {
		log.S().Error(err)
		return nil, err
	}
	if !txnResp.Succeeded {
		resp = (*clientv3.GetResponse)(txnResp.Responses[0].GetResponseRange())
		if resp.Count == 0 {
			return nil, nil
		}
		names = strings.Split(string(resp.Kvs[0].Value), ",")
	}
	realms := make([]*proto.Realm, len(names))
	for i, name := range names {
		r, err := s.GetRealm(ctx, &proto.GetRealmReq{Name: name})
		if err != nil {
			log.S().Errorf("Get Realm %s err %s", name, err)
			return nil, err
		}
		realms[i] = r
	}
	log.S().Infof("locate tenant %s mount point %s successfully %v", req.Tenant, req.MountPoint, realms)
	return &proto.MountPoint{Name: req.MountPoint, Realms: realms}, nil
}

func (s *ZoneServer) getRealmNumClients(realms []string) ([]*realmNumClients, error) {
	var l list
	for _, name := range realms {
		numKey := config.RealmNumClientsPrefix + name
		resp, err := s.eclient.Get(context.TODO(), numKey)
		if err != nil {
			log.S().Errorf("Get Realm %s Num of Clients err %s", name, err)
			return nil, err
		}
		if len(resp.Kvs) > 0 {
			kv := resp.Kvs[0]
			var t = &realmNumClients{}
			t.Realm = string(kv.Key)[len(config.RealmNumClientsPrefix):]
			t.NumClients, _ = strconv.Atoi(string(kv.Value))
			t.ModRevision = kv.ModRevision
			l = append(l, t)
		}
	}
	return l, nil
}

func (s *ZoneServer) RemoveMount(ctx context.Context, req *proto.RemoveMountReq) (*proto.EmptyMessage, error) {
	log.S().Infof("remove tenant %s mount point %s", req.Tenant, req.MountPoint)
	for {
		msg, err := s.removeMount(ctx, req)
		if msg == nil && err == nil {
			continue
		}
		return msg, err
	}
}

func (s *ZoneServer) removeMount(ctx context.Context, req *proto.RemoveMountReq) (*proto.EmptyMessage, error) {
	key := config.MountPrefix + req.Tenant + "/" + req.MountPoint
	strings.ReplaceAll(key, "//", "/")
	resp, err := s.eclient.Get(ctx, key)
	if err != nil {
		log.S().Error(err)
		return nil, err
	}
	if len(resp.Kvs) == 0 || len(resp.Kvs[0].Value) == 0 {
		log.S().Infof("remove tenant %s mount point %s successfully", req.Tenant, req.MountPoint)
		return &proto.EmptyMessage{}, nil
	}
	kv := resp.Kvs[0]
	realmNumClients, err := s.getRealmNumClients(strings.Split(string(kv.Value), ","))
	if err != nil {
		log.S().Error(err)
		return nil, err
	}
	cmps := []clientv3.Cmp{
		clientv3.Compare(clientv3.ModRevision(key), "=", kv.ModRevision),
	}
	ops := []clientv3.Op{
		clientv3.OpDelete(key),
	}
	for _, r := range realmNumClients {
		numKey := config.RealmNumClientsPrefix + r.Realm
		cmp := clientv3.Compare(clientv3.ModRevision(numKey), "=", r.ModRevision)
		cmps = append(cmps, cmp)
		r.NumClients--
		op := clientv3.OpPut(numKey, strconv.Itoa(r.NumClients))
		ops = append(ops, op)
	}
	txn := s.eclient.Txn(context.TODO())
	txn.If(cmps...).
		Then(ops...).
		Else(clientv3.OpGet(key))
	txnResp, err := txn.Commit()
	if err != nil {
		log.S().Error(err)
		return nil, err
	}
	if !txnResp.Succeeded {
		resp = (*clientv3.GetResponse)(txnResp.Responses[0].GetResponseRange())
		if len(resp.Kvs) > 0 && len(resp.Kvs[0].Value) > 0 {
			if kv.ModRevision == resp.Kvs[0].ModRevision {
				return nil, nil
			}
		}
	}
	log.S().Infof("remove tenant %s mount point %s successfully", req.Tenant, req.MountPoint)
	return &proto.EmptyMessage{}, nil
}

func (s *ZoneServer) GetRealm(ctx context.Context, req *proto.GetRealmReq) (*proto.Realm, error) {
	resp, err := s.eclient.Get(ctx, config.RealmMetaPrefix+req.Name)
	if err != nil {
		return nil, err
	}
	if len(resp.Kvs) > 0 {
		kv := resp.Kvs[0]
		var t = &proto.Realm{}
		codec.Unmarshal(kv.Value, t)
		return t, nil
	}
	return nil, errors.New("realm not found")
}

func (s *ZoneServer) GetAllRealms(ctx context.Context, req *proto.EmptyMessage) (*proto.RealmsList, error) {
	log.S().Infof("get all realms")
	resp, err := s.eclient.Get(context.TODO(), config.RealmMetaPrefix, clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	res := &proto.RealmsList{}
	for _, kv := range resp.Kvs {
		var t = &proto.Realm{}
		codec.Unmarshal(kv.Value, t)
		res.Realms = append(res.Realms, t)
	}
	return res, nil
}

func (s *ZoneServer) GlobRealms(ctx context.Context, req *proto.GlobRealmsReq) (*proto.RealmsList, error) {
	log.S().Infof("glob realms")
	p := req.Pattern
	resp, err := s.eclient.Get(context.TODO(), config.RealmMetaPrefix+p, clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	res := &proto.RealmsList{}
	for _, kv := range resp.Kvs {
		var t = &proto.Realm{}
		codec.Unmarshal(kv.Value, t)
		res.Realms = append(res.Realms, t)
	}
	return res, nil
}

func (s *ZoneServer) WatchRealms(req *proto.WatchRealmsReq, srv proto.ZoneService_WatchRealmsServer) error {
	resp, err := s.GlobRealms(context.TODO(), &proto.GlobRealmsReq{Pattern: req.Pattern})
	if err != nil {
		return err
	}
	matches := make([]*realmMeta, len(resp.Realms))
	for i, t := range resp.Realms {
		value, err := codec.Marshal(t)
		if err != nil {
			return err
		}
		matches[i] = s.getOrCreateRealmMeta(value)
	}

	log.S().Infof("watch realms %v, matches %d", resp.Realms, len(matches))

	if len(matches) == 0 {
		return errors.New("no realms")
	}

	ch := make(chan *proto.Realm, 65536)
	pending := make([]*proto.Realm, len(matches))
	for i, m := range matches {
		w := &realmWatcher{ch: ch}
		m.mu.Lock()
		m.mu.watchers[w] = struct{}{}
		pending[i] = cloneRealm(m.Realm)
		m.mu.Unlock()
	}

	h := &watchRealmsHandler{
		ch:     ch,
		stream: srv,
	}
	ctx, cancel := context.WithCancel(h.stream.Context())
	defer cancel()
	go func() {
		for _, p := range pending {
			ch <- p
		}
	}()
	return h.run(ctx)
}

type watchRealmsHandler struct {
	ch     chan *proto.Realm
	stream proto.ZoneService_WatchRealmsServer
}

func (h *watchRealmsHandler) run(ctx context.Context) error {
	resp := &proto.RealmsList{}
	for {
		for i := range resp.Realms {
			resp.Realms[i] = nil
		}
		resp.Realms = resp.Realms[:0]
		select {
		case <-ctx.Done():
			return ctx.Err()
		case z := <-h.ch:
			resp.Realms = append(resp.Realms, z)
		}
		pending := len(h.ch)
		for i := 0; i < pending; i++ {
			z := <-h.ch
			resp.Realms = append(resp.Realms, z)
		}
		if err := h.stream.SendMsg(resp); err != nil {
			return err
		}
	}
}

func (s *ZoneServer) watchThrottlers() {
	resp, err := s.eclient.Get(context.TODO(), config.ThrottlerPrefix, clientv3.WithPrefix())
	if err != nil {
		log.S().Fatalf("failed to get throttlers: %v", err)
	}
	for _, kv := range resp.Kvs {
		s.createThrottlerMeta(kv.Value)
	}
	rev := resp.Header.Revision + 1
	go func() {
		watcher := clientv3.NewWatcher(s.eclient)
		watchChan := watcher.Watch(s.ctx, config.ThrottlerPrefix, clientv3.WithPrefix(), clientv3.WithRev(rev))
		for watchResponse := range watchChan {
			for _, event := range watchResponse.Events {
				switch event.Type {
				case clientv3.EventTypePut:
					s.createThrottlerMeta(event.Kv.Value)
				case clientv3.EventTypeDelete:
					s.deleteThrottlerMeta(string(event.Kv.Key)[len(config.ThrottlerPrefix):])
				}
			}
		}
	}()
}

func (s *ZoneServer) createThrottlerMeta(value []byte) *throttlerMeta {
	var t = &proto.Throttler{}
	codec.Unmarshal(value, t)
	s.throttlers.Lock()
	m, ok := s.throttlers.metas[t.Address]
	if ok {
		changed := string(m.value) != string(value)
		m.Throttler = t
		m.value = value
		clone := cloneThrottler(m.Throttler)
		s.throttlers.Unlock()
		if changed {
			m.broadcast(clone)
		}
	} else {
		m = &throttlerMeta{Throttler: t, value: value}
		m.mu.watchers = make(map[*throttlerWatcher]struct{})
		s.throttlers.metas[t.Address] = m
		s.throttlers.Unlock()
	}
	return m
}

func (s *ZoneServer) deleteThrottlerMeta(n string) {
	s.throttlers.Lock()
	if m, ok := s.throttlers.metas[n]; ok {
		m.Throttler = &proto.Throttler{}
		delete(s.throttlers.metas, n)
		clone := cloneThrottler(m.Throttler)
		s.throttlers.Unlock()
		m.broadcast(clone)
	} else {
		s.throttlers.Unlock()
	}
}

func (s *ZoneServer) ThrottlerHeartbeat(ctx context.Context, req *proto.ThrottlerHeartbeatReq) (*proto.ThrottlerHeartbeatResp, error) {
	r := req.Throttler
	if r == nil {
		return nil, errors.New("empty throttler")
	}
	log.S().Debugf("throttler heartbeat address %s", req.Throttler.Address)
	value, err := codec.Marshal(r)
	if err != nil {
		return nil, err
	}
	if _, err := s.eclient.Put(ctx, config.ThrottlerPrefix+r.Address, string(value)); err != nil {
		return nil, err
	}
	return &proto.ThrottlerHeartbeatResp{}, nil
}