package framework

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	// pprof. don't delete me. make golint happy.
	_ "net/http/pprof"
	"os/signal"
	"runtime/debug"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/etcd/pkg/fileutil"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/prometheus/client_golang/prometheus"
	gozk "github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/pingcap/wormhole/config"
	ether "github.com/pingcap/wormhole/pkg/ether-client"
	"github.com/pingcap/wormhole/pkg/mq/zk"
	"github.com/pingcap/wormhole/pkg/util"
	"github.com/pingcap/wormhole/pkg/whpb"
)

const (
	contractLease            = 80 // in second.  2/3 of ether contractLease.
	workerHeartbeatLease     = 15 // in second
	initKafkaBrokersMaxRetry = 100
	initKafkaBrokersBackOff  = time.Second * 1
	renewJobBackOff          = time.Second * 1
	listenLoopBackOff        = time.Second * 1
	heartbeatBackOff         = time.Second * 1
)

//Worker contains configuration options for worker
type Worker struct {
	cfg *config.Config
	typ whpb.WorkerType
	id  uint64

	heartbeatLease time.Duration
	contractLease  time.Duration
	ec             ether.Client
	zkClient       *zk.Client

	mu struct {
		sync.RWMutex
		handlerPool map[uint64]*Handler
	}

	errCh      chan error
	sc         chan os.Signal
	quit       chan struct{}
	ctx        context.Context
	cancel     context.CancelFunc
	wg         sync.WaitGroup
	workidFile *fileutil.LockedFile
}

//New is the init func for start the work frame
func New(cfg *config.Config, workerType whpb.WorkerType) (*Worker, error) {
	w := &Worker{
		cfg:   cfg,
		typ:   workerType,
		errCh: make(chan error),
		sc:    make(chan os.Signal),
		quit:  make(chan struct{}),
	}
	w.mu.handlerPool = make(map[uint64]*Handler)

	zkClient, err := zk.NewFromConnectionString(w.cfg.ZookeeperUrls, time.Second*5, time.Second*60)
	if err != nil {
		return nil, errors.Trace(err)
	}
	w.zkClient = zkClient

	err = initKafkaBrokers(zkClient, w.cfg)
	if err != nil {
		return nil, errors.Trace(err)
	}
	log.Infof("[worker] get config: %+v", w.cfg)

	w.ec, err = ether.NewClient(w.cfg.EtherEndpoints, workerType, w.cfg.StatusPort)
	if err != nil {
		return nil, errors.Trace(err)
	}
	err = w.initID(w.cfg.WorkerIDFile)
	if err != nil {
		return nil, errors.Trace(err)
	}

	if cfg.HeartbeatLease == 0 {
		w.heartbeatLease = time.Duration(workerHeartbeatLease) * time.Second
	} else {
		w.heartbeatLease = time.Duration(cfg.HeartbeatLease) * time.Second
	}
	if cfg.ContractLease == 0 {
		w.contractLease = time.Duration(contractLease) * time.Second
	} else {
		w.contractLease = time.Duration(cfg.ContractLease) * time.Second
	}

	w.ctx, w.cancel = context.WithCancel(context.Background())

	w.wg.Add(3)
	go w.heartbeatLoop()
	go w.listenLoop()
	go w.renewLoop()

	go w.prometheusHandle()
	go w.signalHandle()

	return w, nil
}

func initKafkaBrokers(zkClient *zk.Client, cfg *config.Config) (err error) {
	var brokerMap map[int32]string
	for i := 0; i < initKafkaBrokersMaxRetry; i++ {
		if i != 0 {
			time.Sleep(initKafkaBrokersBackOff)
			log.Infof("[worker] retry %d to get kafka brokers", i)
		}

		brokerMap, err = zkClient.Brokers()
		if err != nil {
			if util.OriginError(err) == gozk.ErrNoNode || util.OriginError(err) == gozk.ErrNoServer {
				time.Sleep(initKafkaBrokersBackOff)
				continue
			}
			return errors.Trace(err)
		}
		if len(brokerMap) != 0 {
			break
		}
	}
	if err != nil {
		return errors.Trace(errors.Annotatef(err, "reach max retry %d and still failed", initKafkaBrokersMaxRetry))
	}

	log.Infof("[worker] kafka brokers %+v", brokerMap)
	if len(brokerMap) == 0 {
		log.Warnf("[worker] kafka brokers not found in zookeeper")
		return errors.New("kafka brokers not found in zookeeper")
	}

	brokers := make([]string, 0, len(brokerMap))
	for _, broker := range brokerMap {
		brokers = append(brokers, broker)
	}
	cfg.KafkaBrokers = brokers
	return nil
}

//Run to call the Close() and implement
func (w *Worker) Run() {
	for {
		select {
		case err := <-w.errCh:
			log.Errorf("[worker] catch err:%s", errors.ErrorStack(err))
		case <-w.quit:
			// FIXME? if there still has some err from w.errCh not consumed
			w.close()
			return
		}
	}
}

// close releases all resources binded to manager and can only called by Run().
func (w *Worker) close() {
	log.Infof("[worker]name: %s closing...", w.Name())
	w.cancel()
	w.wg.Wait()

	for _, handler := range w.mu.handlerPool {
		handler.Close()
	}
	if w.ec != nil {
		w.ec.Close()
	}
	if w.zkClient != nil {
		w.zkClient.Close()
	}
	if w.workidFile != nil {
		err := w.workidFile.Close()
		if err != nil {
			log.Errorf("unlock file failed,  %v", err)
		}
	}
	log.Infof("[worker] name: %s closed", w.Name())
}

func (w *Worker) catchError(err error) {
	w.errCh <- err
}

// Name returns name of the worker
func (w *Worker) Name() string {
	return fmt.Sprintf("%s-%d", w.typ.String(), w.id)
}

func (w *Worker) listenLoop() {
	log.Debug("[worker] listenLoop goroutine start")
	defer func() {
		if err := recover(); err != nil {
			log.Errorf("[worker] panic. err: %s, stack: %s", err, debug.Stack())
		}
		w.wg.Done()
	}()

	for {
		select {
		case <-w.ctx.Done():
			log.Infof("[worker] listenLoop quit: %s", w.ctx.Err())
			return
		default:
			hbClient, err := w.ec.WorkerHeartbeatClient(w.ctx)
			if err != nil {
				w.catchError(errors.Annotate(err, "[worker] listenLoop get HbClient failed"))
				time.Sleep(listenLoopBackOff)
				continue
			}
			resp, err := hbClient.Recv()
			if err != nil {
				s, ok := status.FromError(util.OriginError(err))
				if ok {
					if s.Code() == codes.Canceled {
						log.Infof("[worker] listenLoop quit. ether code %v", s.Code())
						w.quit <- struct{}{}
						return
					}
				}
				w.catchError(errors.Annotate(err, "[worker] listenLoop heartbeatClient Recv failed"))
				time.Sleep(listenLoopBackOff)
				continue
			}

			if err := resp.GetHeader().GetError(); err != nil {
				w.catchError(errors.Errorf("code %v message %v", err.GetErrcode(), err.GetMessage()))
				time.Sleep(listenLoopBackOff)
				continue
			}

			log.Infof("[worker] recv jobs %+v", resp)
			w.appendNewJobs(resp.NewJobs)
			w.clearObsoleteJobs(resp.ObsoleteJobIds)
		}
	}
}

func (w *Worker) heartbeatLoop() {
	log.Debug("[worker] heartbeatLoop goroutine start")
	defer func() {
		if err := recover(); err != nil {
			log.Errorf("[worker] panic. err: %s, stack: %s", err, debug.Stack())
		}
		w.wg.Done()
	}()

	timer := time.NewTicker(w.heartbeatLease)
	defer timer.Stop()

	for {
		err := w.sendHeartbeat()
		if err != nil {
			w.catchError(errors.Annotate(err, "[worker] heartbeatLoop send heartbeat failed"))
			time.Sleep(heartbeatBackOff)
			continue
		}

		select {
		case <-timer.C:
		case <-w.ctx.Done():
			log.Infof("[worker] heartbeatLoop quit: %s", w.ctx.Err())
			return
		}
	}
}

func (w *Worker) renewLoop() {
	log.Debug("[worker] renewLoop goroutine start")
	defer func() {
		if err := recover(); err != nil {
			log.Errorf("[worker] panic. err: %s, stack: %s", err, debug.Stack())
		}
		w.wg.Done()
	}()

	timer := time.NewTicker(w.contractLease)
	defer timer.Stop()

	for {
		select {
		case <-timer.C:
			err := w.renewJobs()
			if err != nil {
				w.catchError(errors.Annotate(err, "[worker] renewLoop renew jobs fail"))
				time.Sleep(renewJobBackOff)
				continue
			}

		case <-w.ctx.Done():
			log.Infof("[worker] renewLoop quit: %s", w.ctx.Err())
			return
		}
	}
}

func (w *Worker) prometheusHandle() {
	http.Handle("/metrics", prometheus.Handler())
	log.Infof("[worker] listening on status port %d", w.cfg.StatusPort)
	err := http.ListenAndServe(fmt.Sprintf(":%d", w.cfg.StatusPort), nil)
	if err != nil {
		log.Fatalf("[worker] start prometheus server fail: %v", err)
	}
}

func (w *Worker) signalHandle() {
	signal.Notify(w.sc,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT,
		syscall.SIGABRT)
	sig := <-w.sc
	log.Infof("[worker] Got signal [%v] to exit.", sig)
	w.quit <- struct{}{}
}

func (w *Worker) renewJobs() error {
	w.mu.RLock()
	if len(w.mu.handlerPool) == 0 {
		log.Debug("[worker] renewjobs has nothing to do. handler_pool is empty")
		w.mu.RUnlock()
		return nil
	}

	jobIDs := make([]uint64, 0, len(w.mu.handlerPool))
	for jobID := range w.mu.handlerPool {
		jobIDs = append(jobIDs, jobID)
	}
	w.mu.RUnlock()

	log.Infof("[worker] renewjobs %+v", jobIDs)

	var (
		expiredJobs []uint64
		errorJobs   []uint64
		err         error
	)
	maxRetry := 3
	for i := 0; i < maxRetry; i++ {
		expiredJobs, errorJobs, err = w.ec.RenewJob(context.Background(), jobIDs, int64(w.contractLease))
		if err != nil {
			return errors.Trace(err)
		}
		w.clearExpiredJobs(expiredJobs)
		if len(errorJobs) == 0 {
			return nil
		}
		//retry with errorJobs
		log.Warnf("[worker] renewjobs %v failed and retry %d times", errorJobs, i)
		jobIDs = errorJobs
		time.Sleep(renewJobBackOff)
	}
	log.Errorf("[worker] renewjobs %v with max retry %d and still failed", errorJobs, maxRetry)
	return nil
}

func (w *Worker) initID(fname string) error {

	// get from local file
	f, err := fileutil.TryLockFile(fname, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		if err == fileutil.ErrLocked {
			return errors.Errorf("[worker] workerid file %s is used by other process", fname)
		}
		return errors.Trace(err)
	}
	data, err := ioutil.ReadAll(f)
	if err != nil {
		return errors.Trace(err)
	}
	value := strings.TrimSpace(string(data))
	if value == "" {
		w.id, err = w.ec.GetWorkerID()
		if err != nil {
			return errors.Trace(err)
		}
		_, err = f.WriteString(strconv.FormatUint(w.id, 10))
		if err != nil {
			return errors.Trace(err)
		}
		log.Infof("[worker] register new workerid %d", w.id)
	} else {
		id, err := strconv.ParseUint(value, 10, 64)
		if err != nil {
			return errors.Trace(err)
		}
		w.id = id
		log.Infof("[worker] get existed workerid %d", w.id)
	}

	w.ec.SetWorkerID(w.id)
	w.workidFile = f
	return nil
}

func (w *Worker) appendNewJobs(jobs []*whpb.Job) {
	if len(jobs) == 0 {
		return
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	for _, job := range jobs {
		if _, ok := w.mu.handlerPool[job.Id]; !ok {
			h := NewHandler(w.cfg, w.ec, w.zkClient)
			err := h.AssignJob(job)
			if err != nil {
				w.catchError(errors.Annotate(err, "[worker] append new job fail"))
			} else {
				w.mu.handlerPool[job.Id] = h
			}
		}
	}
}

func (w *Worker) clearObsoleteJobs(jobIDs []uint64) {
	if len(jobIDs) == 0 {
		return
	}
	log.Infof("[worker] clear obsolete jobs %+v", jobIDs)
	w.clearJobs(jobIDs)
}

func (w *Worker) clearExpiredJobs(jobIDs []uint64) {
	if len(jobIDs) == 0 {
		return
	}
	log.Infof("[worker] clear expired jobs %+v", jobIDs)
	w.clearJobs(jobIDs)
}

func (w *Worker) clearResolvedJobs(jobIDs []uint64) {
	if len(jobIDs) == 0 {
		return
	}
	log.Infof("[worker] clear resolved jobs %+v", jobIDs)
	w.clearJobs(jobIDs)
}

func (w *Worker) clearJobs(jobIDs []uint64) {
	w.mu.Lock()
	defer w.mu.Unlock()
	for _, jobID := range jobIDs {
		handler, ok := w.mu.handlerPool[jobID]
		if !ok {
			log.Warnf("[worker] jobid %d not found in the worker manager", jobID)
			continue
		}
		handler.Close()
		delete(w.mu.handlerPool, jobID)
	}
}

func (w *Worker) sendHeartbeat() error {
	statusList := make([]*whpb.JobStatus, 0, len(w.mu.handlerPool))
	w.mu.RLock()
	for _, handler := range w.mu.handlerPool {
		statusList = append(statusList, handler.GetStatus())
	}
	w.mu.RUnlock()

	err := w.ec.SendHeartbeat(w.ctx, statusList)
	if err != nil {
		return errors.Trace(err)
	}
	resolvedJobs := make([]uint64, 0, 0)
	for _, status := range statusList {
		if status.GetStage() == whpb.JobStage_Resolved {
			resolvedJobs = append(resolvedJobs, status.GetId())
		}
	}
	w.clearResolvedJobs(resolvedJobs)
	return nil
}