// Copyright 2020 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package pdapi

import (
	"context"
	"crypto/tls"
	"encoding/binary"
	"fmt"
	"time"

	"github.com/pingcap/kvproto/pkg/keyspacepb"
	etcdclientv3 "go.etcd.io/etcd/clientv3"
	etcdclientv3util "go.etcd.io/etcd/clientv3/clientv3util"
	"k8s.io/klog/v2"
)

const (
	pdClusterIDPath     = "/pd/cluster_id"
	keyspacesMetaPrefix = "keyspaces/meta/"
	ZoneRedundant       = "zone_redundant"
	ServerlessTenantID  = "serverless_tenant_id"
)

type KeyValue struct {
	Key   string
	Value []byte
}

type PDEtcdClient interface {
	// Get the specific kvs.
	// if prefix is true will return all kvs with the specified key as prefix
	Get(key string, prefix bool) (kvs []*KeyValue, err error)
	// PutKey will put key to the target pd etcd cluster
	PutKey(key, value string) error
	// PutKey will put key with ttl to the target pd etcd cluster
	PutTTLKey(key, value string, ttl int64) error
	// DeleteKey will delete key from the target pd etcd cluster
	DeleteKey(key string) error
	// WatchKeyspaces will watch matching keyspace changes.
	WatchKeyspaces(match func(keyspace *KeyspaceMeta) bool, watch WatchKeyspaceMeta) (err error)
	// Close will close the etcd connection
	Close() error
}

type pdEtcdClient struct {
	timeout    time.Duration
	etcdClient *etcdclientv3.Client
}

func NewPdEtcdClient(url string, timeout time.Duration, tlsConfig *tls.Config) (PDEtcdClient, error) {
	etcdClient, err := etcdclientv3.New(etcdclientv3.Config{
		Endpoints:   []string{url},
		DialTimeout: timeout,
		TLS:         tlsConfig,
	})
	if err != nil {
		return nil, err
	}
	return &pdEtcdClient{
		etcdClient: etcdClient,
		timeout:    timeout,
	}, nil
}

func (c *pdEtcdClient) Close() error {
	return c.etcdClient.Close()
}

func (c *pdEtcdClient) Get(key string, prefix bool) (kvs []*KeyValue, err error) {
	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()

	var ops []etcdclientv3.OpOption
	if prefix {
		ops = append(ops, etcdclientv3.WithPrefix())
	}
	resp, err := c.etcdClient.Get(ctx, key, ops...)
	if err != nil {
		return nil, err
	}

	for _, kv := range resp.Kvs {
		kvs = append(kvs, &KeyValue{
			Key:   string(kv.Key),
			Value: kv.Value,
		})
	}

	return
}

func (c *pdEtcdClient) PutKey(key, value string) error {
	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()
	_, err := c.etcdClient.Put(ctx, key, value)
	return err
}

func (c *pdEtcdClient) PutTTLKey(key, value string, ttl int64) error {
	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()
	lease, err := c.etcdClient.Grant(ctx, ttl)
	if err != nil {
		return err
	}
	_, err = c.etcdClient.Put(ctx, key, value, etcdclientv3.WithLease(lease.ID))
	return err
}

func (c *pdEtcdClient) DeleteKey(key string) error {
	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()
	kvc := etcdclientv3.NewKV(c.etcdClient)

	// perform a delete only if key already exists
	_, err := kvc.Txn(ctx).
		If(etcdclientv3util.KeyExists(key)).
		Then(etcdclientv3.OpDelete(key)).
		Commit()
	if err != nil {
		return err
	}
	return nil
}

func (c *pdEtcdClient) WatchKeyspaces(match func(keyspace *KeyspaceMeta) bool, watch WatchKeyspaceMeta) (err error) {
	resp, err := c.etcdClient.Get(context.TODO(), pdClusterIDPath)
	if err != nil || len(resp.Kvs) == 0 {
		return err
	}
	b := resp.Kvs[0].Value
	if len(b) != 8 {
		return nil
	}
	clusterID := binary.BigEndian.Uint64(b)
	keyspacePrefix := fmt.Sprintf("/pd/%d/%s", clusterID, keyspacesMetaPrefix)
	resp, err = c.etcdClient.Get(context.TODO(), keyspacePrefix, etcdclientv3.WithPrefix())
	if err != nil {
		return err
	}
	keyspaces := make([]*KeyspaceMeta, 0, len(resp.Kvs))
	for _, kv := range resp.Kvs {
		keyspace, err := UnmarshalKeyspaceMeta(kv.Value)
		if err != nil {
			return err
		}
		if match == nil || match(keyspace) {
			keyspaces = append(keyspaces, keyspace)
		}
	}
	klog.Infof("WatchKeyspaces: get matched keyspaces %d, all keyspaces %d", len(keyspaces), len(resp.Kvs))
	watch <- keyspaces
	rev := resp.Header.Revision + 1
	watchChan := c.etcdClient.Watch(context.TODO(), keyspacePrefix, etcdclientv3.WithPrefix(), etcdclientv3.WithRev(rev))
	for watchResponse := range watchChan {
		keyspaceMetas := make([]*KeyspaceMeta, 0, len(watchResponse.Events))
		for _, event := range watchResponse.Events {
			switch event.Type {
			case etcdclientv3.EventTypePut:
				keyspace, err := UnmarshalKeyspaceMeta(event.Kv.Value)
				if err != nil {
					return err
				}
				if match == nil || match(keyspace) {
					keyspaceMetas = append(keyspaceMetas, keyspace)
				}
			}
		}
		klog.Infof("WatchKeyspaces: watch matched keyspaces %d, changed keyspaces %d", len(keyspaceMetas), len(watchResponse.Events))
		if len(keyspaceMetas) > 0 {
			watch <- keyspaceMetas
		}
	}
	return nil
}

type WatchKeyspaceMeta chan []*KeyspaceMeta

// KeyspaceMeta wraps keyspacepb.KeyspaceMeta to provide custom JSON marshal.
type KeyspaceMeta struct {
	*keyspacepb.KeyspaceMeta
}

func (meta *KeyspaceMeta) IsZoneRedundant() bool {
	keyspaceMeta := meta.KeyspaceMeta
	if keyspaceMeta.State != keyspacepb.KeyspaceState_ENABLED {
		return false
	}
	if v, ok := keyspaceMeta.Config[ZoneRedundant]; ok && v == "true" {
		return true
	}
	return false
}

func (meta *KeyspaceMeta) Clone() *KeyspaceMeta {
	b, _ := meta.KeyspaceMeta.Marshal()
	clone := &keyspacepb.KeyspaceMeta{}
	clone.Unmarshal(b)
	return &KeyspaceMeta{KeyspaceMeta: clone}
}

func MatchZoneRedundant(meta *KeyspaceMeta) bool {
	zoneRedundant, ok := meta.KeyspaceMeta.Config[ZoneRedundant]
	if ok {
		tenantID := meta.KeyspaceMeta.Config[ServerlessTenantID]
		klog.Infof("MatchZoneRedundant: matched keyspace id %d, name %s, state %s, zone redundant [%s] tenant id [%s]", meta.GetId(), meta.GetName(), meta.GetState().String(), zoneRedundant, tenantID)
	}
	return ok
}

func UnmarshalKeyspaceMeta(data []byte) (meta *KeyspaceMeta, err error) {
	keyspace := &keyspacepb.KeyspaceMeta{}
	if err := keyspace.Unmarshal(data); err != nil {
		return nil, err
	}
	meta = &KeyspaceMeta{}
	meta.KeyspaceMeta = keyspace
	return meta, nil
}