package universe

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

	"github.com/pingcap/log"
	"github.com/tidbcloud/ffs/config"
	"github.com/tidbcloud/ffs/pkg/v1/proto"
	"github.com/tidbcloud/ffs/util/codec"
	"go.etcd.io/etcd/clientv3"
	"go.etcd.io/etcd/etcdserver/api/v3rpc/rpctypes"
)

type leaseMeta struct {
	LeaseID clientv3.LeaseID
	lease   clientv3.Lease
}

func cloneZone(z *proto.Zone) *proto.Zone {
	b, _ := z.Marshal()
	clone := &proto.Zone{}
	clone.Unmarshal(b)
	return clone
}

type zoneMeta struct {
	*proto.Zone
	mu struct {
		sync.RWMutex
		watchers map[*zoneWatcher]struct{}
	}
}

func (m *zoneMeta) broadcast(z *proto.Zone) {
	m.mu.RLock()
	for w := range m.mu.watchers {
		w.notify(z)
	}
	m.mu.RUnlock()
}

type zoneWatcher struct {
	ch chan *proto.Zone
}

func (w *zoneWatcher) notify(z *proto.Zone) {
	w.ch <- z
}

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

type tenantMeta struct {
	*proto.Tenant
	value []byte
	mu    struct {
		sync.RWMutex
		watchers map[*tenantWatcher]struct{}
	}
}

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

type tenantWatcher struct {
	ch chan *proto.Tenant
}

func (w *tenantWatcher) notify(t *proto.Tenant) {
	w.ch <- t
}

type UniverseServer struct {
	proto.UnimplementedUniverseServiceServer
	address     string
	storeClient *clientv3.Client
	ctx         context.Context
	cancel      context.CancelFunc
	leaseMeta   leaseMeta
	zones       struct {
		sync.RWMutex
		metas     map[string]*zoneMeta
		leaseLock sync.RWMutex
		leases    map[string]*leaseMeta
	}
	tenants struct {
		sync.RWMutex
		metas map[string]*tenantMeta
	}
}

func NewUniverseServer(address string, storeClient *clientv3.Client) *UniverseServer {
	svr := &UniverseServer{address: address, storeClient: storeClient}
	svr.zones.metas = make(map[string]*zoneMeta)
	svr.zones.leases = make(map[string]*leaseMeta)
	svr.tenants.metas = make(map[string]*tenantMeta)
	svr.ctx, svr.cancel = context.WithCancel(context.Background())
	svr.watchZones()
	svr.watchTenants()
	go svr.run()
	return svr
}

func (s *UniverseServer) Stop() {
	s.revoke()
	s.cancel()
	s.storeClient.Close()
}

func (s *UniverseServer) run() {
	if err := s.register(); err != nil {
		log.S().Fatalf("failed to register %s err %s", s.address, err)
	}
	log.S().Infof("register universe service")
	ticker := time.NewTicker(time.Second)
	for {
		s.keepAliveOnce()
		select {
		case <-ticker.C:
		case <-s.ctx.Done():
			ticker.Stop()
			log.S().Infof("run exit...")
			return
		}
	}
}

func (s *UniverseServer) register() (err error) {
	s.leaseMeta.LeaseID = 0
	kv := clientv3.NewKV(s.storeClient)
	s.leaseMeta.lease = clientv3.NewLease(s.storeClient)
	leaseResp, err := s.leaseMeta.lease.Grant(context.TODO(), 10)
	if err != nil {
		log.S().Errorf("grant err %s", err)
		return err
	}
	_, err = kv.Put(context.TODO(), config.UniversePrefix+s.address, "", clientv3.WithLease(leaseResp.ID))
	if err != nil {
		log.S().Errorf("put err %s", err)
		return err
	}
	s.leaseMeta.LeaseID = leaseResp.ID
	return
}

func (s *UniverseServer) keepAliveOnce() (err error) {
	_, err = s.leaseMeta.lease.KeepAliveOnce(context.TODO(), s.leaseMeta.LeaseID)
	if err != nil {
		if err == rpctypes.ErrLeaseNotFound {
			s.register()
			err = nil
		}
		return
	}
	//log.S().Debugf("keepalive LeaseID:%d", s.leaseMeta.LeaseID)
	return err
}

func (s *UniverseServer) revoke() (err error) {
	ctx, cancel := context.WithTimeout(context.TODO(), time.Millisecond*500)
	defer cancel()
	_, err = s.storeClient.Revoke(ctx, s.leaseMeta.LeaseID)
	if err != nil {
		return
	}
	log.S().Infof("revoke LeaseID:%d", s.leaseMeta.LeaseID)
	return
}

func (s *UniverseServer) watchZones() {
	resp, err := s.storeClient.Get(context.TODO(), config.ZonePrefix, clientv3.WithPrefix())
	if err != nil {
		log.S().Fatalf("failed to get zones: %v", err)
	}
	for _, kv := range resp.Kvs {
		k := string(kv.Key[len(config.ZonePrefix):])
		ks := strings.Split(k, "/")
		if len(ks) == 2 {
			log.S().Infof("init zone %s endpoint %s", ks[0], ks[1])
			s.getOrCreateZoneMeta(ks[0], ks[1])
		}
	}
	rev := resp.Header.Revision + 1
	go func() {
		watchChan := s.storeClient.Watch(s.ctx, config.ZonePrefix, clientv3.WithPrefix(), clientv3.WithRev(rev))
		for watchResponse := range watchChan {
			for _, event := range watchResponse.Events {
				switch event.Type {
				case clientv3.EventTypePut:
					k := string(event.Kv.Key[len(config.ZonePrefix):])
					ks := strings.Split(k, "/")
					if len(ks) == 2 {
						log.S().Infof("create zone %s endpoint %s", ks[0], ks[1])
						s.getOrCreateZoneMeta(ks[0], ks[1])
					}
				case clientv3.EventTypeDelete:
					k := string(event.Kv.Key[len(config.ZonePrefix):])
					ks := strings.Split(k, "/")
					if len(ks) == 2 {
						log.S().Infof("delete zone %s endpoint %s", ks[0], ks[1])
						s.deleteZoneMeta(ks[0], ks[1])
					}
				}
			}
		}
	}()
}

func (s *UniverseServer) getOrCreateZoneMeta(name string, endpoints ...string) *zoneMeta {
	s.zones.Lock()
	m, ok := s.zones.metas[name]
	if ok {
		var changed = false
		for i := range endpoints {
			endpoint := endpoints[i]
			existed := exists(m.Endpoints, endpoint)
			if !existed {
				m.Endpoints = append(m.Endpoints, endpoint)
			}
			changed = changed || !existed
		}
		clone := cloneZone(m.Zone)
		s.zones.Unlock()
		if changed {
			m.broadcast(clone)
		}
	} else {
		m = &zoneMeta{Zone: &proto.Zone{Name: name, Endpoints: endpoints}}
		m.mu.watchers = make(map[*zoneWatcher]struct{})
		s.zones.metas[name] = m
		s.zones.Unlock()
	}
	return m
}

func (s *UniverseServer) deleteZoneMeta(name, endpoint string) *zoneMeta {
	s.zones.Lock()
	m, ok := s.zones.metas[name]
	if ok {
		if len(endpoint) > 0 {
			m.Zone.Endpoints, _ = remove(m.Zone.Endpoints, endpoint)
		} else {
			delete(s.zones.metas, name)
		}
		clone := cloneZone(m.Zone)
		s.zones.Unlock()
		m.broadcast(clone)
	} else {
		s.zones.Unlock()
	}
	return m
}

func (s *UniverseServer) ZoneHeartbeat(ctx context.Context, req *proto.ZoneHeartbeatReq) (*proto.Zone, error) {
	z := req.Zone
	if z != nil && len(z.Endpoints) > 0 {
		endpoint := z.Endpoints[0]
		//log.S().Debugf("zone heartbeat: name %s address %s", z.Name, endpoint)
		key := config.ZonePrefix + z.Name + "/" + endpoint
		for {
			s.zones.leaseLock.RLock()
			lm, ok := s.zones.leases[key]
			s.zones.leaseLock.RUnlock()
			if !ok {
				lease := clientv3.NewLease(s.storeClient)
				//ttl := int64(config.Timeout / time.Second)
				leaseResp, err := lease.Grant(context.TODO(), 10)
				if err != nil {
					log.S().Error(err)
					return nil, err
				}
				lm = &leaseMeta{LeaseID: leaseResp.ID, lease: lease}
				s.zones.leaseLock.Lock()
				l, ok := s.zones.leases[key]
				if !ok {
					s.zones.leases[key] = lm
				} else {
					lm = l
				}
				s.zones.leaseLock.Unlock()
				if !ok {
					_, err := s.storeClient.Put(context.TODO(), key, "", clientv3.WithLease(lm.LeaseID))
					if err != nil {
						log.S().Error(err)
						s.zones.leaseLock.Lock()
						delete(s.zones.leases, key)
						s.zones.leaseLock.Unlock()
						lm.lease.Revoke(context.TODO(), lm.LeaseID)
						lm.lease.Close()
						time.Sleep(time.Second)
						continue
					}
				} else {
					lease.Revoke(context.TODO(), leaseResp.ID)
					lease.Close()
				}
			}
			if _, err := lm.lease.KeepAliveOnce(context.TODO(), lm.LeaseID); err == rpctypes.ErrLeaseNotFound {
				log.S().Error(err)
				s.zones.leaseLock.Lock()
				delete(s.zones.leases, key)
				s.zones.leaseLock.Unlock()
				lm.lease.Revoke(context.TODO(), lm.LeaseID)
				lm.lease.Close()
				continue
			} else if err != nil {
				log.S().Error(err)
				return nil, err
			}
			break
		}
	}
	return req.Zone, nil
}

func (s *UniverseServer) GetZone(ctx context.Context, req *proto.GetZoneReq) (*proto.Zone, error) {
	resp, err := s.storeClient.Get(ctx, config.ZonePrefix+req.Name+"/", clientv3.WithPrefix())
	if err != nil {
		log.S().Errorf("get zonename %s err %s", req.Name, err)
		return nil, err
	}
	z := &proto.Zone{Name: req.Name}
	for _, kv := range resp.Kvs {
		k := string(kv.Key[len(config.ZonePrefix):])
		ks := strings.Split(k, "/")
		if len(ks) == 2 && ks[0] == req.Name {
			z.Endpoints = append(z.Endpoints, ks[1])
		}
	}
	log.S().Infof("get zonename %s address %v", req.Name, z.Endpoints)
	return z, nil
}

func (s *UniverseServer) DeleteZone(ctx context.Context, req *proto.DeleteZoneReq) (*proto.EmptyMessage, error) {
	n := req.Name
	if len(n) > 0 {
		_, err := s.storeClient.Delete(ctx, config.ZonePrefix+n, clientv3.WithPrefix())
		if err != nil {
			return nil, err
		}
	}
	return &proto.EmptyMessage{}, nil
}

func (s *UniverseServer) GetAllZones(ctx context.Context, req *proto.EmptyMessage) (*proto.ZonesList, error) {
	resp, err := s.storeClient.Get(context.TODO(), config.ZonePrefix, clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	zones := make(map[string][]string)
	for _, kv := range resp.Kvs {
		k := string(kv.Key[len(config.ZonePrefix):])
		ks := strings.Split(k, "/")
		if len(ks) == 2 {
			name := ks[0]
			endpoint := ks[1]
			if l, ok := zones[name]; ok {
				l = append(l, endpoint)
			} else {
				zones[name] = []string{endpoint}
			}
		}
	}
	res := &proto.ZonesList{}
	for k, v := range zones {
		res.Zones = append(res.Zones, &proto.Zone{Name: k, Endpoints: v})
	}
	return res, nil
}

func (s *UniverseServer) GlobZones(ctx context.Context, req *proto.GlobZonesReq) (*proto.ZonesList, error) {
	p := req.Pattern
	resp, err := s.storeClient.Get(context.TODO(), config.ZonePrefix+p, clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	zones := make(map[string][]string)
	for _, kv := range resp.Kvs {
		k := string(kv.Key[len(config.ZonePrefix):])
		ks := strings.Split(k, "/")
		if len(ks) == 2 {
			name := ks[0]
			endpoint := ks[1]
			if l, ok := zones[name]; ok {
				l = append(l, endpoint)
			} else {
				zones[name] = []string{endpoint}
			}
		}
	}
	res := &proto.ZonesList{}
	for k, v := range zones {
		res.Zones = append(res.Zones, &proto.Zone{Name: k, Endpoints: v})
	}
	return res, nil
}

func (s *UniverseServer) WatchZones(req *proto.WatchZonesReq, srv proto.UniverseService_WatchZonesServer) error {
	resp, err := s.GlobZones(context.TODO(), &proto.GlobZonesReq{Pattern: req.Pattern})
	if err != nil {
		return err
	}
	matches := make([]*zoneMeta, len(resp.Zones))
	for i, z := range resp.Zones {
		matches[i] = s.getOrCreateZoneMeta(z.Name, z.Endpoints...)
	}

	log.S().Infof("watch zones %v, matches %d", resp.Zones, len(matches))

	if len(matches) == 0 {
		return errors.New("no zones")
	}
	ch := make(chan *proto.Zone, 65536)
	pending := make([]*proto.Zone, len(matches))
	for i, m := range matches {
		w := &zoneWatcher{ch: ch}
		m.mu.Lock()
		m.mu.watchers[w] = struct{}{}
		pending[i] = cloneZone(m.Zone)
		m.mu.Unlock()
	}
	h := &watchZonesHandler{
		ch:     ch,
		stream: srv,
	}
	ctx, cancel := context.WithCancel(h.stream.Context())
	defer cancel()
	go func() {
		for _, z := range pending {
			ch <- z
		}
	}()
	return h.run(ctx)
}

type watchZonesHandler struct {
	ch     chan *proto.Zone
	stream proto.UniverseService_WatchZonesServer
}

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

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

func (s *UniverseServer) getOrCreateTenantMeta(value []byte) *tenantMeta {
	var t = &proto.Tenant{}
	codec.Unmarshal(value, t)
	s.tenants.Lock()
	m, ok := s.tenants.metas[t.Name]
	if ok {
		changed := string(m.value) != string(value)
		m.Tenant = t
		m.value = value
		clone := cloneTenant(m.Tenant)
		s.tenants.Unlock()
		if changed {
			m.broadcast(clone)
		}
	} else {
		m = &tenantMeta{Tenant: t, value: value}
		m.mu.watchers = make(map[*tenantWatcher]struct{})
		s.tenants.metas[t.Name] = m
		s.tenants.Unlock()
	}
	return m
}

func (s *UniverseServer) deleteTenantMeta(n string) {
	s.tenants.Lock()
	if m, ok := s.tenants.metas[n]; ok {
		m.Tenant = &proto.Tenant{}
		delete(s.tenants.metas, n)
		clone := cloneTenant(m.Tenant)
		s.tenants.Unlock()
		m.broadcast(clone)
	} else {
		s.tenants.Unlock()
	}
}

func (s *UniverseServer) CreateTenant(ctx context.Context, req *proto.CreateTenantReq) (*proto.Tenant, error) {
	t := req.Tenant
	if t == nil {
		return nil, errors.New("empty tenant")
	}
	log.S().Infof("create tenant name: %s", t.Name)
	value, err := codec.Marshal(t)
	if err != nil {
		return nil, err
	}
	if _, err := s.storeClient.Put(ctx, config.TenantPrefix+t.Name, string(value)); err != nil {
		return nil, err
	}
	return t, nil
}

func (s *UniverseServer) GetTenant(ctx context.Context, req *proto.GetTenantReq) (*proto.Tenant, error) {
	resp, err := s.storeClient.Get(ctx, config.TenantPrefix+req.Name)
	if err != nil {
		return nil, err
	}
	if len(resp.Kvs) > 0 {
		kv := resp.Kvs[0]
		var t = &proto.Tenant{}
		codec.Unmarshal(kv.Value, t)
		return t, nil
	}
	return nil, errors.New("tenant not found")
}

func (s *UniverseServer) DeleteTenant(ctx context.Context, req *proto.DeleteTenantReq) (*proto.EmptyMessage, error) {
	n := req.Name
	if len(n) > 0 {
		_, err := s.storeClient.Delete(ctx, config.TenantPrefix+n)
		if err != nil {
			return nil, err
		}
	}
	return &proto.EmptyMessage{}, nil
}

func (s *UniverseServer) GetAllTenants(ctx context.Context, req *proto.EmptyMessage) (*proto.TenantsList, error) {
	resp, err := s.storeClient.Get(context.TODO(), config.TenantPrefix, clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	res := &proto.TenantsList{}
	for _, kv := range resp.Kvs {
		var t = &proto.Tenant{}
		codec.Unmarshal(kv.Value, t)
		res.Tenants = append(res.Tenants, t)
	}
	return res, nil
}

func (s *UniverseServer) GlobTenants(ctx context.Context, req *proto.GlobTenantsReq) (*proto.TenantsList, error) {
	p := req.Pattern
	resp, err := s.storeClient.Get(context.TODO(), config.TenantPrefix+p, clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	res := &proto.TenantsList{}
	for _, kv := range resp.Kvs {
		var t = &proto.Tenant{}
		codec.Unmarshal(kv.Value, t)
		res.Tenants = append(res.Tenants, t)
	}
	return res, nil
}

func (s *UniverseServer) WatchTenants(req *proto.WatchTenantsReq, srv proto.UniverseService_WatchTenantsServer) error {
	resp, err := s.GlobTenants(context.TODO(), &proto.GlobTenantsReq{Pattern: req.Pattern})
	if err != nil {
		return err
	}
	matches := make([]*tenantMeta, len(resp.Tenants))
	for i, t := range resp.Tenants {
		value, err := codec.Marshal(t)
		if err != nil {
			return err
		}
		matches[i] = s.getOrCreateTenantMeta(value)
	}

	log.S().Infof("watch tenants %v, matches %d", resp.Tenants, len(matches))

	if len(matches) == 0 {
		return errors.New("no tenants")
	}
	ch := make(chan *proto.Tenant, 65536)
	pending := make([]*proto.Tenant, len(matches))
	for i, m := range matches {
		w := &tenantWatcher{ch: ch}
		m.mu.Lock()
		m.mu.watchers[w] = struct{}{}
		pending[i] = cloneTenant(m.Tenant)
		m.mu.Unlock()
	}

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

type watchTenantsHandler struct {
	ch     chan *proto.Tenant
	stream proto.UniverseService_WatchTenantsServer
}

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

func exists(l []string, addr string) bool {
	for i := range l {
		if l[i] == addr {
			return true
		}
	}
	return false
}

func remove(s []string, addr string) ([]string, bool) {
	for i := range s {
		if s[i] == addr {
			s[i] = s[len(s)-1]
			return s[:len(s)-1], true
		}
	}
	return s, false
}