// Copyright 2023 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"context"
	"encoding/json"
	"net/http"
	"sync"
	"syscall"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/tidb/pkg/config"
	"github.com/pingcap/tidb/pkg/domain"
	"github.com/pingcap/tidb/pkg/kv"
	"github.com/pingcap/tidb/pkg/parser/terror"
	"github.com/pingcap/tidb/pkg/session"
	storeerr "github.com/pingcap/tidb/pkg/store/driver/error"
	"github.com/pingcap/tidb/pkg/util"
	"github.com/pingcap/tidb/pkg/util/logutil"
	"github.com/pingcap/tidb/pkg/util/signal"
	"github.com/pingcap/tidb/pkg/util/sqlexec"
	"go.uber.org/zap"
)

// HealthHandler is the http handler for health check.
type HealthHandler struct {
	dom *domain.Domain
	sm  util.SessionManager

	// Fields for periodic health check
	mu                 sync.RWMutex
	checkCtx           context.Context
	checkCancel        context.CancelFunc
	lastCheckStatus    string
	lastCheckErr       error
	lastSuccessTime    time.Time
	isPeriodicChecking bool
}

// NewHealthHandler creates a new health handler.
func NewHealthHandler(dom *domain.Domain, sm util.SessionManager) *HealthHandler {
	return &HealthHandler{dom: dom, sm: sm}
}

type healthResponse struct {
	Status string `json:"status"`
	Token  string `json:"token"`
	Error  string `json:"error,omitempty"`
}

// checkOnce performs a single health check and returns status and error.
// Status is "up" or "down", error is set if the check fails (excluding throttled errors).
func (h *HealthHandler) checkOnce(ctx context.Context) (status string, err error) {
	ctx = &deadlineToCanceledCtx{ctx}

	// Use defer to handle throttled errors uniformly
	defer func() {
		status = "up"
		if err != nil {
			if terror.ErrorEqual(err, storeerr.ErrResourceGroupThrottled) {
				// Ignore throttled error
				err = nil
				return
			}
			status = "down"
			logutil.BgLogger().Warn("health check failed", zap.Error(err))
		}
	}()

	if h.dom == nil {
		return // it may be nil in test.
	}

	se, err := session.CreateSessionWithDomain(h.dom.Store(), h.dom)
	if se != nil {
		defer se.Close()
	}
	if err != nil {
		return
	}
	se.SetSessionManager(h.sm)
	ctx = kv.WithInternalSourceType(ctx, kv.InternalTxnOthers)
	rs, err := se.ExecuteInternal(ctx, "SELECT variable_name FROM mysql.tidb LIMIT 1")
	if err != nil {
		return
	}
	defer rs.Close()
	rows, err := sqlexec.DrainRecordSet(ctx, rs, 1)
	if err != nil {
		return
	}
	if len(rows) != 1 {
		err = errors.New("timeout to read mysql.tidb")
	}
	return
}

// StartPeriodicCheck starts periodic health checks using the interval from config.
// If a periodic check is already running, it will be stopped first.
func (h *HealthHandler) StartPeriodicCheck() {
	h.mu.Lock()
	defer h.mu.Unlock()

	// Stop existing check if any
	if h.checkCancel != nil {
		h.checkCancel()
		h.checkCancel = nil
		h.isPeriodicChecking = false
	}

	cfg := config.GetGlobalConfig()
	if cfg.HealthCheckInterval == 0 {
		return
	}
	interval := time.Duration(cfg.HealthCheckInterval) * time.Second

	ctx, cancel := context.WithCancel(context.Background())
	h.checkCtx = ctx
	h.checkCancel = cancel
	h.isPeriodicChecking = true

	go h.periodicCheck(interval)
}

// StopPeriodicCheck stops the periodic health check.
func (h *HealthHandler) StopPeriodicCheck() {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.checkCancel != nil {
		h.checkCancel()
		h.checkCancel = nil
		h.isPeriodicChecking = false
	}
}

// GetLastCheckResult returns the last health check result and timestamp.
func (h *HealthHandler) GetLastCheckResult() (string, error) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.lastCheckStatus, h.lastCheckErr
}

// periodicCheck performs health checks periodically until the context is cancelled.
func (h *HealthHandler) periodicCheck(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	// Capture the context at the start of the goroutine to avoid race conditions
	h.mu.RLock()
	ctx := h.checkCtx
	h.mu.RUnlock()

	// Perform initial check immediately
	h.performCheck(ctx)

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			h.performCheck(ctx)
		}
	}
}

// performCheck performs a single health check with the given context.
func (h *HealthHandler) performCheck(baseCtx context.Context) {
	ctx, cancel := context.WithTimeout(baseCtx, 5*time.Second)
	defer cancel()

	status, err := h.checkOnce(ctx)
	var lastSuccess time.Time

	h.mu.Lock()
	h.lastCheckStatus = status
	h.lastCheckErr = err
	lastSuccess = h.lastSuccessTime
	if status == "up" {
		h.lastSuccessTime = time.Now()
	}
	h.mu.Unlock()

	if status == "up" {
		return
	}

	cfg := config.GetGlobalConfig()
	failureTimeout := time.Duration(cfg.HealthCheckFailureTimeout) * time.Second

	if failureTimeout > 0 && !lastSuccess.IsZero() && time.Since(lastSuccess) > failureTimeout {
		errorMsg := ""
		if err != nil {
			errorMsg = err.Error()
		}
		logutil.BgLogger().Error("health check failed for too long, exiting",
			zap.Duration("failure-duration", time.Since(lastSuccess)),
			zap.Duration("timeout", failureTimeout),
			zap.String("last-error", errorMsg),
			zap.Time("last-success-time", lastSuccess),
		)
		// Try to set force shutdown if sm is a Server instance
		if svr, ok := h.sm.(*Server); ok {
			svr.SetForceShutdown()
		}
		// Consider the server is going to force shutdown, send a high priority signal to kill tidb.
		signal.TiDBExit(syscall.SIGINT)
	}
}

// If we never had a successful check, don't exit yet
// Wait for at least one successful check before starting the failure timer

// ServeHTTP implements the http.Handler interface.
func (h *HealthHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	response := &healthResponse{
		Token: config.GetGlobalKeyspaceName(),
	}

	h.mu.RLock()
	isPeriodicChecking := h.isPeriodicChecking
	lastStatus := h.lastCheckStatus
	lastErr := h.lastCheckErr
	h.mu.RUnlock()

	var status string
	var err error
	if isPeriodicChecking && lastStatus != "" {
		// Use last check result if periodic check is running and has a result
		status = lastStatus
		err = lastErr
	} else {
		// Perform a new check if periodic check is not running or has no result yet
		ctx, cancel := context.WithTimeout(req.Context(), 5*time.Second)
		defer cancel()
		status, err = h.checkOnce(ctx)
	}

	response.Status = status
	if err != nil {
		response.Error = err.Error()
	}
	json.NewEncoder(w).Encode(response)
}

type deadlineToCanceledCtx struct {
	context.Context
}

func (c *deadlineToCanceledCtx) Err() error {
	err := c.Context.Err()
	if err == context.DeadlineExceeded {
		return context.Canceled
	}
	return err
}