package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"sync"
	"time"

	orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
)

// loadPayload 的固定内容同时是处理器的输入和预期结果；不含凭据或外部地址。
type loadPayload struct {
	RunID      string `json:"run_id"`      // RunID 隔离本次运行，不作为业务幂等保证。
	Sequence   int    `json:"sequence"`    // Sequence 从0开始的请求序号，用于发现重复执行及错配。
	DurationMS int64  `json:"duration_ms"` // DurationMS 模拟等待毫秒数，按本次固定时长数组取模。
}

func decodePayload(raw json.RawMessage) (loadPayload, bool) {
	var wire struct {
		RunID      *string `json:"run_id"`
		Sequence   *int    `json:"sequence"`
		DurationMS *int64  `json:"duration_ms"`
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.DisallowUnknownFields()
	if decoder.Decode(&wire) != nil || wire.RunID == nil || wire.Sequence == nil || wire.DurationMS == nil {
		return loadPayload{}, false
	}
	var trailing any
	if decoder.Decode(&trailing) != io.EOF {
		return loadPayload{}, false
	}
	return loadPayload{*wire.RunID, *wire.Sequence, *wire.DurationMS}, true
}

type admission struct {
	sequence int
	task     orbit.Task
	duration time.Duration
	code     string
}

type taskAPI interface {
	SubmitTask(context.Context, string, orbit.SubmitTaskRequest) (orbit.SubmitTaskResult, error)
	GetTask(context.Context, string) (orbit.Task, error)
	ListTaskAttempts(context.Context, string, orbit.HistoryPageRequest) (orbit.AttemptPage, error)
}

type measurement struct {
	queue, execution, completed time.Duration
}

func expected(cfg config, runID string, sequence int) loadPayload {
	return loadPayload{runID, sequence, cfg.Durations[sequence%len(cfg.Durations)].Milliseconds()}
}

func handlerName(runID string) string { return "load.exec." + runID }
func businessKey(runID string) string { return "orbit-exec-" + runID }
func requestKey(runID string, sequence int) string {
	return fmt.Sprintf("orbit-exec-%s-%06d", runID, sequence)
}

// invocationLog 有界保存本次每个序号的实际处理器调用，发现重复时不覆盖首条证据。
type invocationLog struct {
	mu      sync.Mutex
	calls   []int
	taskIDs []string
	invalid int
}

func newInvocationLog(n int) *invocationLog {
	return &invocationLog{calls: make([]int, n), taskIDs: make([]string, n)}
}

func (log *invocationLog) handler(runCtx context.Context, cfg config, runID string) orbit.WorkerHandler {
	return orbit.WorkerHandler{Capability: orbit.WorkerCapability{Handler: handlerName(runID), HandlerVersion: "v1", PayloadVersion: "v1"}, Handle: func(ctx context.Context, execution orbit.WorkerExecution) (json.RawMessage, error) {
		payload, valid := decodePayload(execution.Task.Payload)
		if !valid || payload.RunID != runID || payload.Sequence < 0 || payload.Sequence >= cfg.Requests || payload != expected(cfg, runID, payload.Sequence) || execution.Task.QueueID != cfg.QueueID || execution.Task.BusinessKey != businessKey(runID) || execution.Task.Handler != handlerName(runID) {
			log.mu.Lock()
			log.invalid++
			log.mu.Unlock()
			return nil, &orbit.WorkerFailure{Code: "LOAD_INPUT_MISMATCH"}
		}
		log.mu.Lock()
		log.calls[payload.Sequence]++
		first := log.calls[payload.Sequence] == 1
		if first {
			log.taskIDs[payload.Sequence] = execution.Task.ID
		}
		log.mu.Unlock()
		if !first {
			return nil, &orbit.WorkerFailure{Code: "LOAD_DUPLICATE_EXECUTION"}
		}
		timer := time.NewTimer(time.Duration(payload.DurationMS) * time.Millisecond)
		defer timer.Stop()
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-runCtx.Done():
			return nil, runCtx.Err()
		case <-timer.C:
			return json.Marshal(payload)
		}
	}}
}

func measure(parent context.Context, cfg config, client *orbit.Client, runID string) executionReport {
	r := newReport(cfg, runID)
	ctx, cancel := context.WithTimeout(parent, cfg.RunTimeout)
	defer cancel()
	requestCtx, stop := context.WithTimeout(ctx, cfg.Timeout)
	queue, err := client.GetQueue(requestCtx, cfg.QueueID)
	stop()
	if err != nil {
		r.Errors["PREFLIGHT_"+classify(err)]++
		return r.finish()
	}
	r.Queue = &queue
	if queue.ID != cfg.QueueID || queue.Paused || queue.MaxConcurrency == 0 || queue.RatePerSecond == 0 {
		r.Errors["QUEUE_UNAVAILABLE"]++
		return r.finish()
	}
	log := newInvocationLog(cfg.Requests)
	var workers []*orbit.Worker
	for i := 0; i < cfg.Workers; i++ {
		w, err := orbit.NewWorker(client, orbit.WorkerOptions{QueueID: cfg.QueueID, Capacity: uint32(cfg.Capacity), ShutdownTimeout: 5 * time.Second, Handlers: []orbit.WorkerHandler{log.handler(ctx, cfg, runID)}})
		if err != nil {
			r.Errors["WORKER_CONFIG"]++
			return r.finish()
		}
		workers = append(workers, w)
	}
	workerCtx, stopWorkers := context.WithCancel(ctx)
	done := make(chan error, len(workers))
	for _, w := range workers {
		go func(w *orbit.Worker) {
			err := w.Run(workerCtx)
			if workerCtx.Err() == nil {
				cancel()
			} // 不可恢复连接错误立即停止继续提交。
			done <- err
		}(w)
	}
	admissions := submit(ctx, cfg, client, runID)
	var successful []admission
	var admissionDurations []time.Duration
	seen := map[string]bool{}
	for _, a := range admissions {
		admissionDurations = append(admissionDurations, a.duration)
		if a.code != "" {
			r.AdmissionErrors++
			r.Errors["ADMISSION_"+a.code]++
			continue
		}
		if seen[a.task.ID] {
			r.AdmissionErrors++
			r.Errors["DUPLICATE_TASK_ID"]++
			continue
		}
		seen[a.task.ID] = true
		successful = append(successful, a)
	}
	r.Attempted = len(admissions)
	r.Accepted = len(successful)
	r.NotSubmitted = cfg.Requests - r.Attempted
	r.Admission = distributionOf(admissionDurations)
	r.AdmissionFinishedAt = time.Now().UTC()
	verified := verifyAll(ctx, cfg, client, runID, successful)
	var queued, executed, completed []time.Duration
	for _, v := range verified {
		if v.code != "" {
			r.Errors[v.code]++
			continue
		}
		r.VerifiedSuccess++
		queued = append(queued, v.measurement.queue)
		executed = append(executed, v.measurement.execution)
		completed = append(completed, v.measurement.completed)
	}
	if ctx.Err() != nil {
		r.Errors["RUN_"+classify(ctx.Err())]++
	}
	stopWorkers()
	for range workers {
		if err := <-done; err != nil {
			r.Errors["WORKER_STOP"]++
		}
	}
	for _, w := range workers {
		s := w.Stats()
		r.WorkerStarted += s.Started
		r.Reconnects += s.Reconnects
		r.Renewals += s.RenewalsConfirmed
	}
	log.mu.Lock()
	for i, calls := range log.calls {
		r.HandlerCalls += calls
		if calls > 1 {
			r.DuplicateInvocations += calls - 1
		}
		if calls > 0 && !seen[log.taskIDs[i]] {
			r.UnconfirmedHandlerCalls += calls
		}
	}
	for _, a := range successful {
		if log.calls[a.sequence] != 1 || log.taskIDs[a.sequence] != a.task.ID {
			r.Errors["INVOCATION_MISMATCH"]++
		}
	}
	r.Errors["INVALID_HANDLER_INPUT"] += log.invalid
	log.mu.Unlock()
	r.QueueWait = distributionOf(queued)
	r.AuthorizedExecution = distributionOf(executed)
	r.TaskCompletion = distributionOf(completed)
	return r.finish()
}

func submit(ctx context.Context, cfg config, api taskAPI, runID string) []admission {
	jobs := make(chan int)
	out := make(chan admission, cfg.Requests)
	var wg sync.WaitGroup
	for i := 0; i < cfg.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for sequence := range jobs {
				if ctx.Err() != nil {
					continue
				}
				payload, _ := json.Marshal(expected(cfg, runID, sequence))
				start := time.Now()
				call, stop := context.WithTimeout(ctx, cfg.Timeout)
				res, err := api.SubmitTask(call, requestKey(runID, sequence), orbit.SubmitTaskRequest{QueueID: cfg.QueueID, Handler: handlerName(runID), HandlerVersion: "v1", PayloadVersion: "v1", BusinessKey: businessKey(runID), Payload: payload, TaskTimeoutMS: uint64(cfg.RunTimeout.Milliseconds()), RecoverySafety: orbit.RecoveryReconcileOnly, Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: uint64(cfg.RunTimeout.Milliseconds())}})
				stop()
				a := admission{sequence: sequence, task: res.Task, duration: time.Since(start)}
				if err != nil {
					a.code = classify(err)
				} else if res.Deduplicated || !validID(res.Task.ID) || !matches(res.Task, cfg, runID, sequence, false) {
					a.code = "INVALID_RECEIPT"
				}
				out <- a
			}
		}()
	}
feed:
	for i := 0; i < cfg.Requests; i++ {
		select {
		case <-ctx.Done():
			break feed
		case jobs <- i:
		}
	}
	close(jobs)
	wg.Wait()
	close(out)
	results := make([]admission, 0, cfg.Requests)
	for a := range out {
		results = append(results, a)
	}
	return results
}

type verification struct {
	measurement measurement
	code        string
}

func verifyAll(ctx context.Context, cfg config, api taskAPI, runID string, admissions []admission) []verification {
	jobs := make(chan admission, len(admissions))
	for _, a := range admissions {
		jobs <- a
	}
	close(jobs)
	out := make(chan verification, len(admissions))
	var wg sync.WaitGroup
	for i := 0; i < cfg.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for a := range jobs {
				m, code := verify(ctx, cfg, api, runID, a)
				out <- verification{m, code}
			}
		}()
	}
	wg.Wait()
	close(out)
	results := make([]verification, 0, len(admissions))
	for v := range out {
		results = append(results, v)
	}
	return results
}

func verify(ctx context.Context, cfg config, api taskAPI, runID string, a admission) (measurement, string) {
	for {
		if ctx.Err() != nil {
			return measurement{}, "VERIFY_" + classify(ctx.Err())
		}
		call, stop := context.WithTimeout(ctx, cfg.Timeout)
		task, err := api.GetTask(call, a.task.ID)
		stop()
		if err != nil {
			return measurement{}, "READ_" + classify(err)
		}
		if task.ID != a.task.ID || !matches(task, cfg, runID, a.sequence, false) {
			return measurement{}, "TASK_MISMATCH"
		}
		switch task.State {
		case orbit.TaskSucceeded:
			call, stop = context.WithTimeout(ctx, cfg.Timeout)
			history, err := api.ListTaskAttempts(call, task.ID, orbit.HistoryPageRequest{Limit: 2})
			stop()
			if err != nil {
				return measurement{}, "HISTORY_" + classify(err)
			}
			return checkCompletion(cfg, runID, a.sequence, task, history)
		case orbit.TaskFailed, orbit.TaskCancelled, orbit.TaskNeedsReview:
			return measurement{}, "TASK_" + string(task.State)
		case orbit.TaskWaiting, orbit.TaskLeased, orbit.TaskRunning:
		default:
			return measurement{}, "UNKNOWN_TASK_STATE"
		}
		timer := time.NewTimer(cfg.PollInterval)
		select {
		case <-ctx.Done():
			timer.Stop()
			return measurement{}, "VERIFY_" + classify(ctx.Err())
		case <-timer.C:
		}
	}
}

func matches(task orbit.Task, cfg config, runID string, sequence int, result bool) bool {
	payload, valid := decodePayload(task.Payload)
	if task.QueueID != cfg.QueueID || task.Handler != handlerName(runID) || task.HandlerVersion != "v1" || task.PayloadVersion != "v1" || task.BusinessKey != businessKey(runID) || task.Retry.MaxAttempts != 1 || task.RecoverySafety != orbit.RecoveryReconcileOnly || !valid || payload != expected(cfg, runID, sequence) {
		return false
	}
	if result {
		output, valid := decodePayload(task.Result)
		if !valid || output != expected(cfg, runID, sequence) {
			return false
		}
	}
	return true
}

func checkCompletion(cfg config, runID string, sequence int, task orbit.Task, history orbit.AttemptPage) (measurement, string) {
	if task.State != orbit.TaskSucceeded || task.EffectUncertain || task.AttemptCount != 1 || task.CompletedAt == nil || !matches(task, cfg, runID, sequence, true) {
		return measurement{}, "COMPLETION_MISMATCH"
	}
	if len(history.Items) != 1 || history.NextCursor != "" {
		return measurement{}, "ATTEMPT_COUNT_MISMATCH"
	}
	a := history.Items[0]
	if !validID(a.ID) || a.TaskID != task.ID || a.Epoch != 1 || a.State != orbit.AttemptSucceeded || a.AuthorizedAt == nil || a.EndedAt == nil {
		return measurement{}, "ATTEMPT_MISMATCH"
	}
	if task.CreatedAt.IsZero() || task.AvailableAt.Before(task.CreatedAt) || a.CreatedAt.Before(task.AvailableAt) || a.AuthorizedAt.Before(a.CreatedAt) || a.EndedAt.Before(*a.AuthorizedAt) || !a.EndedAt.Equal(*task.CompletedAt) {
		return measurement{}, "TIMESTAMP_MISMATCH"
	}
	return measurement{queue: a.AuthorizedAt.Sub(task.AvailableAt), execution: a.EndedAt.Sub(*a.AuthorizedAt), completed: task.CompletedAt.Sub(task.CreatedAt)}, ""
}