// Code mostly adapted from pd-client

package ether

import (
	"fmt"
	"net"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/wormhole/pkg/util"
	"github.com/pingcap/wormhole/pkg/whpb"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

const (
	etherTimeout          = 3 * time.Second
	maxInitClusterRetries = 100
)

var (
	errFailInitClusterID = errors.New("[ether-client] failed to get cluster id")
)

// Client is a ether client that should be implemented.
type Client interface {
	// GetClusterID returns the cluster ID from ether.
	GetClusterID() uint64
	// GetWorkerID returns the worker ID from ether.
	GetWorkerID() (uint64, error)
	// SetWorkerID sets worker ID
	SetWorkerID(uint64)
	// GetStorageRootPath returns storage root path in ether.
	GetStorageRootPath() string
	// WorkerHeartbeatClient represents heartbeat interface interacted with ether server.
	WorkerHeartbeatClient(ctx context.Context) (whpb.Ether_WorkerHeartbeatClient, error)
	// SendHeartbeat sends job status to ether.
	SendHeartbeat(ctx context.Context, jobStatuses []*whpb.JobStatus) error
	// RenewJob renew job lease.
	RenewJob(ctx context.Context, jobIDList []uint64, lease int64) (expiredJobs []uint64, errorJobs []uint64, err error)
	// CreateFence creates fence in ether, and receives a fenceid.
	CreateFence(ctx context.Context, topics []*whpb.Topic) (uint64, error)
	// WaitOnFence waits on it's fences util all consumers reach all fence
	WaitOnFenceMidway(fenceID uint64, jobID uint64, topic *whpb.Topic, partitions []int32) (whpb.Ether_WaitOnFenceMidwayClient, error)
	WaitOnFenceFinal(fenceID uint64, jobID uint64, topic *whpb.Topic, partitions []int32) (whpb.Ether_WaitOnFenceFinalClient, error)
	// Close closes the client.
	Close()
}

type client struct {
	urls            []string
	clusterID       uint64
	workerID        uint64
	storageRootPath string
	mu              sync.Mutex

	connMu struct {
		sync.RWMutex
		clientConns map[string]*grpc.ClientConn
		leader      string
	}

	checkLeaderCh chan struct{}

	wg                    sync.WaitGroup
	ctx                   context.Context
	cancel                context.CancelFunc
	workerType            whpb.WorkerType
	localPort             uint64
	cachedHeartbeatClient whpb.Ether_WorkerHeartbeatClient
	cancelHeartbeatClient context.CancelFunc
}

// NewClient creates a ether client.
func NewClient(etherAddrs []string, wt whpb.WorkerType, localPort uint64) (Client, error) {
	log.Infof("[ether-client] create ether client with endpoints %v", etherAddrs)
	ctx, cancel := context.WithCancel(context.Background())
	c := &client{
		urls:          addrsToUrls(etherAddrs),
		checkLeaderCh: make(chan struct{}, 1),
		ctx:           ctx,
		cancel:        cancel,
		workerType:    wt,
		localPort:     localPort,
	}
	c.connMu.clientConns = make(map[string]*grpc.ClientConn)

	if err := c.initClusterID(); err != nil {
		return nil, errors.Trace(err)
	}
	if err := c.updateLeader(); err != nil {
		return nil, errors.Trace(err)
	}

	log.Infof("[ether-client] init cluster id %v", c.clusterID)

	go c.leaderLoop()
	return c, nil
}

func (c *client) initClusterID() error {
	ctx, cancel := context.WithCancel(c.ctx)
	defer cancel()
	for i := 0; i < maxInitClusterRetries; i++ {
		for _, u := range c.urls {
			members, err := c.getMembers(ctx, u)
			if err != nil || members.ClusterId == 0 {
				log.Warnf("[ether-client] failed to get cluster id: %v", err)
				continue
			}
			c.clusterID = members.ClusterId
			c.storageRootPath = members.StorageRootPath
			return nil
		}

		time.Sleep(time.Second)
	}

	return errors.Trace(errFailInitClusterID)
}

func (c *client) updateLeader() error {
	ctx, cancel := context.WithCancel(c.ctx)
	defer cancel()
	for _, u := range c.urls {
		members, err := c.getMembers(ctx, u)
		if err != nil || members.GetLeader() == nil || len(members.GetLeader().ClientUrls) == 0 {
			continue
		}
		if err = c.switchLeader(members.GetLeader().ClientUrls); err != nil {
			return errors.Trace(err)
		}
		return nil
	}
	return errors.Errorf("[ether-client] failed to get leader from %v", c.urls)
}

func (c *client) getMembers(ctx context.Context, url string) (*whpb.GetClusterInfoResponse, error) {
	cc, err := c.getOrCreateGRPCConn(url)
	if err != nil {
		return nil, errors.Trace(err)
	}
	members, err := whpb.NewEtherClient(cc).GetClusterInfo(ctx, &whpb.GetClusterInfoRequest{})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return members, nil
}

func (c *client) switchLeader(addrs []string) error {
	// FIXME: How to safely compare leader urls? For now, only allows one client url.
	addr := addrs[0]

	c.connMu.RLock()
	oldLeader := c.connMu.leader
	c.connMu.RUnlock()

	if addr == oldLeader {
		return nil
	}

	log.Infof("[ether-client] leader switches to: %v, previous: %v", addr, oldLeader)
	if _, err := c.getOrCreateGRPCConn(addr); err != nil {
		return errors.Trace(err)
	}

	c.connMu.Lock()
	defer c.connMu.Unlock()
	c.connMu.leader = addr
	return nil
}

func (c *client) getOrCreateGRPCConn(addr string) (*grpc.ClientConn, error) {
	c.connMu.RLock()
	conn, ok := c.connMu.clientConns[addr]
	c.connMu.RUnlock()
	if ok {
		return conn, nil
	}

	cc, err := grpc.Dial(addr, grpc.WithDialer(func(addr string, d time.Duration) (net.Conn, error) {
		u, err := url.Parse(addr)
		if err != nil {
			return nil, errors.Trace(err)
		}
		// For tests.
		if u.Scheme == "unix" || u.Scheme == "unixs" {
			return net.DialTimeout(u.Scheme, u.Host, d)
		}
		return net.DialTimeout("tcp", u.Host, d)
	}), grpc.WithInsecure()) // TODO: Support HTTPS.
	if err != nil {
		return nil, errors.Trace(err)
	}

	c.connMu.Lock()
	defer c.connMu.Unlock()
	if old, ok := c.connMu.clientConns[addr]; ok {
		cc.Close()
		return old, nil
	}

	c.connMu.clientConns[addr] = cc
	return cc, nil
}

func (c *client) leaderLoop() {
	c.wg.Add(1)
	defer c.wg.Done()

	ctx, cancel := context.WithCancel(c.ctx)
	defer cancel()

	for {
		select {
		case <-c.checkLeaderCh:
		case <-time.After(time.Minute):
		case <-ctx.Done():
			return
		}

		if err := c.updateLeader(); err != nil {
			log.Errorf("[ether-client] failed updateLeader: %v", err)
		}
	}
}

func (c *client) GetClusterID() uint64 {
	return c.clusterID
}

func (c *client) GetWorkerID() (uint64, error) {
	if c.workerID != 0 {
		return c.workerID, nil
	}
	// lazy register.
	err := c.registerWorker(c.ctx, c.workerType)
	if err != nil {
		return 0, errors.Trace(err)
	}
	return c.workerID, nil
}

func (c *client) SetWorkerID(workerID uint64) {
	c.workerID = workerID
}

func (c *client) GetStorageRootPath() string {
	return c.storageRootPath
}

func (c *client) WorkerHeartbeatClient(ctx context.Context) (whpb.Ether_WorkerHeartbeatClient, error) {
	return c.getOrCreateHeartbeatClient(ctx)
}

func (c *client) getOrCreateHeartbeatClient(ctx context.Context) (whpb.Ether_WorkerHeartbeatClient, error) {
	// the function may be accessed concurrently in framework/worker.go. add mutex lock here.
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.cachedHeartbeatClient != nil {
		// send empty heartbeat message to check connection ok or not.
		err := c.cachedHeartbeatClient.Send(&whpb.WorkerHeartbeatRequest{
			Header: c.requestHeader(),
		})
		if err != nil {
			log.Warnf("[ether-client] send empty heartbeat failed %s", err.Error())
			c.cancelHeartbeatClient()
			goto NewHeartbeatClient
		}
		return c.cachedHeartbeatClient, nil
	}

NewHeartbeatClient:
	ctx, cancel := context.WithCancel(ctx)
	log.Info("[ether-client] create new worker heartbeat client")
	hbClient, err := c.leaderClient().WorkerHeartbeat(ctx)
	if err != nil {
		cancel()
		return nil, errors.Trace(err)
	}
	// cache it for later use.
	c.cachedHeartbeatClient = hbClient
	c.cancelHeartbeatClient = cancel
	return hbClient, nil
}

func (c *client) SendHeartbeat(ctx context.Context, jobStatuses []*whpb.JobStatus) error {
	client, err := c.getOrCreateHeartbeatClient(ctx)
	if err != nil {
		return errors.Trace(err)
	}
	req := &whpb.WorkerHeartbeatRequest{
		Header: c.requestHeader(),
		Jobs:   jobStatuses,
	}
	return errors.Trace(client.Send(req))
}

func (c *client) RenewJob(ctx context.Context, jobIDList []uint64, lease int64) (expiredJobs []uint64, errorJobs []uint64, err error) {
	ctx, cancel := context.WithTimeout(ctx, etherTimeout)
	resp, err := c.leaderClient().RenewJob(ctx, &whpb.RenewJobRequest{
		Header: c.requestHeader(),
		JobIds: jobIDList,
		Lease:  lease,
	})
	cancel()

	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	if err := resp.GetHeader().GetError(); err != nil {
		if err.Errcode == whpb.ErrorCode_NotLeader {
			c.scheduleCheckLeader()
		}
		return nil, nil, errors.New(err.String())
	}
	return resp.ExpiredJobs, resp.ErrorJobs, nil
}

func (c *client) CreateFence(ctx context.Context, topics []*whpb.Topic) (uint64, error) {
	ctx, cancel := context.WithTimeout(ctx, etherTimeout)
	resp, err := c.leaderClient().CreateFence(ctx, &whpb.CreateFenceRequest{
		Header: c.requestHeader(),
		Topics: topics,
	})
	cancel()

	if err != nil {
		return 0, errors.Trace(err)
	}
	if err := resp.GetHeader().GetError(); err != nil {
		if err.Errcode == whpb.ErrorCode_NotLeader {
			c.scheduleCheckLeader()
		}
		return 0, errors.New(err.String())
	}
	return resp.GetFenceId(), nil
}

// WaitOnFence waits on it's fences util all consumers reach all fences.
func (c *client) WaitOnFenceMidway(fenceID uint64, jobID uint64, topic *whpb.Topic, partitions []int32) (whpb.Ether_WaitOnFenceMidwayClient, error) {
	ctx, cancel := context.WithCancel(c.ctx)
	client, err := c.leaderClient().WaitOnFenceMidway(ctx, &whpb.WaitOnFenceRequest{
		Header:     c.requestHeader(),
		FenceId:    fenceID,
		JobId:      jobID,
		Topic:      topic,
		Partitions: partitions,
	})

	if err != nil {
		cancel()
		return nil, errors.Trace(err)
	}
	return client, nil
}

func (c *client) WaitOnFenceFinal(fenceID uint64, jobID uint64, topic *whpb.Topic, partitions []int32) (whpb.Ether_WaitOnFenceFinalClient, error) {
	ctx, cancel := context.WithCancel(c.ctx)
	client, err := c.leaderClient().WaitOnFenceFinal(ctx, &whpb.WaitOnFenceRequest{
		Header:     c.requestHeader(),
		FenceId:    fenceID,
		JobId:      jobID,
		Topic:      topic,
		Partitions: partitions,
	})

	if err != nil {
		cancel()
		return nil, errors.Trace(err)
	}
	return client, nil
}

func (c *client) registerWorker(ctx context.Context, wt whpb.WorkerType) error {
	host, err := util.GetLocalIP()
	if err != nil {
		return errors.Trace(err)
	}
	ctx, cancel := context.WithTimeout(ctx, etherTimeout)
	req := &whpb.RegisterWorkerRequest{
		Header: c.requestHeader(),
		Worker: &whpb.Worker{
			Type: wt,
			Annotations: map[string]string{
				"host": fmt.Sprintf("%s:%d", host, c.localPort), // TODO: rename
			},
			CreateTime: time.Now().Unix(),
			LastTime:   time.Now().Unix(),
		},
	}
	log.Debugf("[ether-client] register worker req %+v", req)
	resp, err := c.leaderClient().RegisterWorker(ctx, req)
	cancel()
	if err != nil {
		return errors.Trace(err)
	}
	if err := resp.Header.GetError(); err != nil {
		if err.Errcode == whpb.ErrorCode_NotLeader {
			c.scheduleCheckLeader()
		}
		return errors.New(err.String())
	}

	c.workerID = resp.GetHeader().GetWorkerId()
	log.Debugf("[ether-client] register worker, got workerid %d", c.workerID)
	return nil
}

func (c *client) Close() {
	c.cancel()
	c.wg.Wait()

	c.connMu.Lock()
	defer c.connMu.Unlock()
	for _, cc := range c.connMu.clientConns {
		if err := cc.Close(); err != nil {
			log.Errorf("[ether-client] failed close grpc clientConn: %v", err)
		}
	}
}

func (c *client) leaderClient() whpb.EtherClient {
	c.connMu.RLock()
	defer c.connMu.RUnlock()

	return whpb.NewEtherClient(c.connMu.clientConns[c.connMu.leader])
}

func (c *client) scheduleCheckLeader() {
	select {
	case c.checkLeaderCh <- struct{}{}:
	default:
	}
}

func (c *client) requestHeader() *whpb.RequestHeader {
	return &whpb.RequestHeader{
		ClusterId: c.clusterID,
		WorkerId:  c.workerID,
	}
}

func addrsToUrls(addrs []string) []string {
	// Add default schema "http://" to addrs.
	urls := make([]string, 0, len(addrs))
	for _, addr := range addrs {
		if strings.Contains(addr, "://") {
			urls = append(urls, addr)
		} else {
			urls = append(urls, "http://"+addr)
		}
	}
	return urls
}