package executor

import (
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	runningWalltimeResultLabel       = "running"
	defaultWalltimeUpdateMinInterval = 5 * time.Second
)

var (
	MetricsEventCountTotal = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "runtime",
			Subsystem: "executor",
			Name:      "event_count_total",
			Help:      "Executor events counter",
		}, []string{"executor_name", "event_name"})

	MetricsDagActionDurationHistogram = prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Namespace: "runtime",
			Subsystem: "executor",
			Name:      "dag_action_duration_seconds",
			Help:      "Executor dag action elapsed seconds histogram",
			Buckets:   prometheus.ExponentialBuckets(0.001, 2, 24), // 1ms ~ 2.3h
		}, []string{"executor_name", "task_name", "action_name", "result"})

	MetricsDagActionRunningWalltimeSecsGauge = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "runtime",
			Subsystem: "executor",
			Name:      "dag_action_running_walltime_seconds",
			Help:      "Executor dag action running walltime in seconds",
		}, []string{"executor_name", "task_name", "action_name", "result"})
)

func init() {
	prometheus.MustRegister(MetricsEventCountTotal)
	prometheus.MustRegister(MetricsDagActionDurationHistogram)
	prometheus.MustRegister(MetricsDagActionRunningWalltimeSecsGauge)
}

type MetricsRecorder struct {
	executorName            string
	locker                  sync.Mutex
	running                 map[string]runningActionStatus
	lastWalltimeUpdateAt    time.Time
	walltimeUpdateMinPeriod time.Duration
}

type runningActionStatus struct {
	taskName         string
	actionName       string
	accumulated      time.Duration
	currentAttemptAt time.Time
}

func NewMetricsRecorder(executorName string) *MetricsRecorder {
	return &MetricsRecorder{
		executorName:            executorName,
		running:                 make(map[string]runningActionStatus),
		walltimeUpdateMinPeriod: defaultWalltimeUpdateMinInterval,
	}
}

func (recorder *MetricsRecorder) OnEvent(name EventName, data ...EventData) {
	MetricsEventCountTotal.WithLabelValues(recorder.executorName, string(name)).Inc()

	switch name {
	case EventActionTaskStart:
		actionTask := data[0].Object.(*ActionTask)
		recorder.locker.Lock()
		actionName := actionTask.DagAction.Action.Name()
		taskName := actionTask.ExecTask.ExecId
		actionTaskID := actionTask.Id()
		status, ok := recorder.running[actionTaskID]
		if !ok {
			status.taskName = taskName
			status.actionName = actionName
		}
		status.currentAttemptAt = time.Now()
		recorder.running[actionTaskID] = status
		MetricsDagActionRunningWalltimeSecsGauge.
			WithLabelValues(recorder.executorName, status.taskName, status.actionName, runningWalltimeResultLabel).
			Set(status.accumulated.Seconds())
		recorder.locker.Unlock()
	case EventExecutorRunLoopStart:
		now := time.Now()
		recorder.locker.Lock()
		if len(recorder.running) == 0 {
			recorder.locker.Unlock()
			return
		}
		if now.Sub(recorder.lastWalltimeUpdateAt) < recorder.walltimeUpdateMinPeriod {
			recorder.locker.Unlock()
			return
		}
		for _, running := range recorder.running {
			walltime := running.accumulated
			if !running.currentAttemptAt.IsZero() {
				walltime += now.Sub(running.currentAttemptAt)
			}
			MetricsDagActionRunningWalltimeSecsGauge.
				WithLabelValues(recorder.executorName, running.taskName, running.actionName, runningWalltimeResultLabel).
				Set(walltime.Seconds())
		}
		recorder.lastWalltimeUpdateAt = now
		recorder.locker.Unlock()
	case EventExecutorRunLoopOnActionResult:
		actionTaskStatus := data[0].Object.(*ActionTaskStatus)
		actionName := actionTaskStatus.ActionTask.DagAction.Action.Name()
		taskName := actionTaskStatus.ActionTask.ExecTask.ExecId
		actionTaskID := actionTaskStatus.ActionTask.Id()

		recorder.locker.Lock()
		running, ok := recorder.running[actionTaskID]
		if ok {
			// Trust executor-reported elapsed when start time exists, including zero duration.
			if !actionTaskStatus.StartTime.IsZero() {
				running.accumulated += actionTaskStatus.Elapsed
			} else if !running.currentAttemptAt.IsZero() {
				// Fallback only when elapsed is unavailable from executor.
				running.accumulated += time.Since(running.currentAttemptAt)
			}
			running.currentAttemptAt = time.Time{}
		}
		if ok && actionTaskStatus.Error != nil {
			recorder.running[actionTaskID] = running
		} else {
			delete(recorder.running, actionTaskID)
		}
		recorder.locker.Unlock()

		if ok && actionTaskStatus.Error == nil {
			MetricsDagActionRunningWalltimeSecsGauge.
				DeleteLabelValues(recorder.executorName, running.taskName, running.actionName, runningWalltimeResultLabel)
		} else if ok {
			MetricsDagActionRunningWalltimeSecsGauge.
				WithLabelValues(recorder.executorName, running.taskName, running.actionName, runningWalltimeResultLabel).
				Set(running.accumulated.Seconds())
		}

		if actionTaskStatus.Error == nil {
			MetricsDagActionDurationHistogram.
				WithLabelValues(recorder.executorName, taskName, actionName, "ok").
				Observe(actionTaskStatus.Elapsed.Seconds())
		} else {
			MetricsDagActionDurationHistogram.
				WithLabelValues(recorder.executorName, taskName, actionName, "error").
				Observe(actionTaskStatus.Elapsed.Seconds())
		}
	case EventExecTaskDone:
		execTask := data[0].Object.(*ExecutorTask)
		recorder.locker.Lock()
		for actionTaskID, running := range recorder.running {
			if running.taskName != execTask.ExecId {
				continue
			}
			MetricsDagActionRunningWalltimeSecsGauge.
				DeleteLabelValues(recorder.executorName, running.taskName, running.actionName, runningWalltimeResultLabel)
			delete(recorder.running, actionTaskID)
		}
		recorder.locker.Unlock()
	case EventExecutorDone:
		recorder.locker.Lock()
		for _, running := range recorder.running {
			MetricsDagActionRunningWalltimeSecsGauge.
				DeleteLabelValues(recorder.executorName, running.taskName, running.actionName, runningWalltimeResultLabel)
		}
		clear(recorder.running)
		recorder.locker.Unlock()
	}
}