package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"sync"
	"time"

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

// schedulePayload 冻结运行和计划序号；精确字段集合用于检测错配输入。
type schedulePayload struct {
	RunID    string `json:"run_id"`   // RunID 本次测量身份，不是业务幂等保证。
	Schedule int    `json:"schedule"` // Schedule 从零开始的计划序号。
}

func decodePayload(raw json.RawMessage) (schedulePayload, bool) {
	// map中的RawMessage保留整数文本，不接受缺字段、null和非精确字段名。
	var fields map[string]json.RawMessage
	var out schedulePayload
	if json.Unmarshal(raw, &fields) != nil || len(fields) != 2 {
		return out, false
	}
	run, hasRun := fields["run_id"]
	index, hasIndex := fields["schedule"]
	if !hasRun || !hasIndex || string(run) == "null" || string(index) == "null" || json.Unmarshal(run, &out.RunID) != nil || json.Unmarshal(index, &out.Schedule) != nil {
		return schedulePayload{}, false
	}
	// 按原文计数，拒绝重复键；字段次序可以不同。
	var keys []string
	decoder := json.NewDecoder(bytes.NewReader(raw))
	if _, err := decoder.Token(); err != nil {
		return schedulePayload{}, false
	}
	for decoder.More() {
		k, err := decoder.Token()
		if err != nil {
			return schedulePayload{}, false
		}
		keys = append(keys, k.(string))
		var value json.RawMessage
		if decoder.Decode(&value) != nil {
			return schedulePayload{}, false
		}
	}
	return out, len(keys) == 2
}

func handlerName(runID string) string { return "load.schedule." + runID }
func operationKey(runID string, index int, op string) string {
	return fmt.Sprintf("schedule-%s-%02d-%s", runID, index, op)
}
func businessKey(runID string, index int) string { return operationKey(runID, index, "task") }

type scheduleAPI interface {
	GetQueue(context.Context, string) (orbit.Queue, error)
	CreateSchedule(context.Context, string, orbit.CreateScheduleRequest) (orbit.ScheduleResult, error)
	GetSchedule(context.Context, string) (orbit.Schedule, error)
	ResumeSchedule(context.Context, string, string, orbit.ControlScheduleRequest) (orbit.ScheduleResult, error)
	PauseSchedule(context.Context, string, string, orbit.ControlScheduleRequest) (orbit.ScheduleResult, error)
	ListScheduleOccurrences(context.Context, string, orbit.SchedulePageRequest) (orbit.ScheduleOccurrencePage, error)
	GetTask(context.Context, string) (orbit.Task, error)
	ListTaskAttempts(context.Context, string, orbit.HistoryPageRequest) (orbit.AttemptPage, error)
	GetTaskSource(context.Context, string) (orbit.TaskSource, error)
}

func scheduleDefinition(cfg config, runID string, index int, start time.Time) orbit.ScheduleDefinition {
	end := start.Add(time.Duration(cfg.Requests) * cfg.Interval)
	payload, _ := json.Marshal(schedulePayload{RunID: runID, Schedule: index})
	return orbit.ScheduleDefinition{
		Rule:   orbit.ScheduleRule{Kind: "FIXED_RATE", Timezone: "UTC", StartAt: start, EndAt: &end, IntervalMS: uint64(cfg.Interval.Milliseconds())},
		Policy: orbit.SchedulePolicy{Misfire: "CATCH_UP", Overlap: "PARALLEL", CatchupLimit: 100, CatchupWindowMS: uint64((cfg.RunTimeout + time.Duration(cfg.Requests)*cfg.Interval + cfg.StartDelay).Milliseconds()), GraceMS: 1000},
		Task:   orbit.SubmitTaskRequest{QueueID: cfg.QueueID, Handler: handlerName(runID), HandlerVersion: "v1", PayloadVersion: "v1", BusinessKey: businessKey(runID, index), Payload: payload, TaskTimeoutMS: uint64(cfg.RunTimeout.Milliseconds()), RecoverySafety: orbit.RecoveryReconcileOnly, Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: uint64(cfg.RunTimeout.Milliseconds())}},
	}
}

func measure(parent context.Context, cfg config, api *orbit.Client, runID string) (r scheduleReport) {
	r = newReport(cfg, runID)
	ctx, cancel := context.WithTimeout(parent, cfg.RunTimeout)
	defer cancel()
	var workers []*orbit.Worker
	var stopWorkers context.CancelFunc
	var done chan error
	log := newInvocationLog(cfg.Requests * cfg.Schedules)
	observations := newObservations(cfg, runID)
	// 所有返回路径都执行有界收尾；超时后的暂停使用独立上下文，结果不明不换键重试。
	defer func() {
		if stopWorkers != nil {
			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
		}
		cleanup, stop := context.WithTimeout(context.Background(), 5*time.Second)
		var wg sync.WaitGroup
		codes := make(chan string, len(r.Plans))
		for i := range r.Plans {
			wg.Add(1)
			go func(p *planReport) {
				defer wg.Done()
				if validID(p.ID) {
					codes <- pauseSchedule(cleanup, cfg, api, runID, p)
				}
			}(&r.Plans[i])
		}
		wg.Wait()
		stop()
		close(codes)
		for code := range codes {
			if code != "" {
				r.Errors[code]++
			}
		}
		observations.summarize(&r)
		log.summarize(&r, observations)
		if ctx.Err() != nil {
			r.Errors["RUN_"+classify(ctx.Err())]++
		}
		r = r.finish()
	}()
	call, stop := context.WithTimeout(ctx, cfg.Timeout)
	queue, err := api.GetQueue(call, cfg.QueueID)
	stop()
	if err != nil {
		r.Errors["PREFLIGHT_"+classify(err)]++
		return
	}
	r.Queue = &queue
	if queue.ID != cfg.QueueID || !validID(queue.NamespaceID) || queue.Paused || queue.MaxConcurrency == 0 || queue.RatePerSecond == 0 {
		r.Errors["QUEUE_UNAVAILABLE"]++
		return
	}
	r.ScheduleStart = time.Now().UTC().Truncate(time.Microsecond).Add(cfg.StartDelay)
	r.ScheduleEnd = r.ScheduleStart.Add(time.Duration(cfg.Requests) * cfg.Interval)
	// 在恢复任何计划之前完成全部创建；计划有限终止，即使进程被强杀也不会无限触发。
	seen := map[string]bool{}
	for i := 0; i < cfg.Schedules; i++ {
		if ctx.Err() != nil {
			return
		}
		definition := scheduleDefinition(cfg, runID, i, r.ScheduleStart)
		key := operationKey(runID, i, "create")
		p := planReport{Index: i, CreateKey: key}
		r.Plans = append(r.Plans, p)
		call, stop = context.WithTimeout(ctx, cfg.Timeout)
		res, e := api.CreateSchedule(call, key, orbit.CreateScheduleRequest{Name: key, Paused: true, Definition: definition})
		stop()
		if e != nil {
			r.Errors["CREATE_"+classify(e)]++
			return
		}
		p.ID = res.Schedule.ID
		p.CalendarVersion = res.Schedule.Current.Definition.Rule.CalendarVersion
		p.EffectiveAt = res.Schedule.Current.EffectiveAt
		r.Plans[i] = p
		if res.Deduplicated || !validID(p.ID) || seen[p.ID] || !matchesSchedule(res.Schedule, cfg, runID, i, queue.NamespaceID, r.ScheduleStart) || !res.Schedule.Paused || res.Schedule.Current.EffectiveAt.After(r.ScheduleStart) {
			r.Errors["CREATE_MISMATCH"]++
			return
		}
		seen[p.ID] = true
		observations.addPlan(p.ID, i, r.ScheduleStart, queue.NamespaceID)
	}
	for i := 0; i < cfg.Workers; i++ {
		w, e := orbit.NewWorker(api, orbit.WorkerOptions{QueueID: cfg.QueueID, Capacity: uint32(cfg.Capacity), ShutdownTimeout: 5 * time.Second, Handlers: []orbit.WorkerHandler{log.handler(ctx, cfg, runID)}})
		if e != nil {
			r.Errors["WORKER_CONFIG"]++
			return
		}
		workers = append(workers, w)
	}
	workerCtx, workerCancel := context.WithCancel(ctx)
	stopWorkers = workerCancel
	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)
	}
	for i := range r.Plans {
		p := &r.Plans[i]
		call, stop = context.WithTimeout(ctx, cfg.Timeout)
		current, e := api.GetSchedule(call, p.ID)
		stop()
		if e != nil {
			r.Errors["RESUME_READ_"+classify(e)]++
			return
		}
		if !time.Now().Before(r.ScheduleStart) || current.ID != p.ID || !matchesSchedule(current, cfg, runID, i, queue.NamespaceID, r.ScheduleStart) || !current.Paused || !sameVersion(current, *p) {
			r.Errors["SETUP_LATE_OR_CHANGED"]++
			return
		}
		p.ResumeKey = operationKey(runID, i, "resume")
		p.ResumeRevision = current.Revision
		call, stop = context.WithTimeout(ctx, cfg.Timeout)
		res, e := api.ResumeSchedule(call, p.ID, p.ResumeKey, orbit.ControlScheduleRequest{ExpectedRevision: current.Revision, Reason: "start bounded schedule verification"})
		stop()
		if e != nil {
			r.Errors["RESUME_"+classify(e)]++
			return
		}
		if res.Deduplicated || res.Schedule.ID != p.ID || res.Schedule.Paused || !matchesSchedule(res.Schedule, cfg, runID, i, queue.NamespaceID, r.ScheduleStart) || !res.Schedule.UpdatedAt.Before(r.ScheduleStart) || !sameVersion(res.Schedule, *p) {
			r.Errors["RESUME_MISMATCH_OR_LATE"]++
			return
		}
		p.Resumed = true
	}
	for {
		if ctx.Err() != nil {
			return
		}
		all := true
		for i := range r.Plans {
			p := &r.Plans[i]
			call, stop = context.WithTimeout(ctx, cfg.Timeout)
			page, e := api.ListScheduleOccurrences(call, p.ID, orbit.SchedulePageRequest{Limit: 100})
			stop()
			if e != nil {
				r.Errors["OCCURRENCES_"+classify(e)]++
				return
			}
			if code := observations.ingest(p.ID, page); code != "" {
				r.Errors[code]++
				return
			}
			complete, code := observations.verify(ctx, api, p.ID)
			if code != "" {
				r.Errors[code]++
				return
			}
			if complete {
				call, stop = context.WithTimeout(ctx, cfg.Timeout)
				current, e := api.GetSchedule(call, p.ID)
				stop()
				if e != nil {
					r.Errors["FINAL_READ_"+classify(e)]++
					return
				}
				if current.ID != p.ID || !matchesSchedule(current, cfg, runID, i, queue.NamespaceID, r.ScheduleStart) || current.Paused || !sameVersion(current, *p) {
					r.Errors["FINAL_SCHEDULE_MISMATCH"]++
					return
				}
				p.Exhausted = current.CheckAt == nil && current.NextFireAt == nil && current.ProcessingVersion == 1
			}
			all = all && complete && p.Exhausted
		}
		if all {
			return
		}
		timer := time.NewTimer(cfg.PollInterval)
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case <-timer.C:
		}
	}
}

func matchesSchedule(s orbit.Schedule, cfg config, runID string, index int, namespace string, start time.Time) bool {
	want := scheduleDefinition(cfg, runID, index, start)
	got := s.Current.Definition
	// 服务端冻结具体时区快照标识；不把API版本v1误当作时区数据版本。
	if !validCalendarVersion(got.Rule.CalendarVersion) {
		return false
	}
	got.Rule.CalendarVersion = ""
	return s.NamespaceID == namespace && s.QueueID == cfg.QueueID && s.Name == operationKey(runID, index, "create") && s.CurrentVersion == 1 && s.Current.Version == 1 && s.Revision > 0 && !s.Current.EffectiveAt.IsZero() && sameDefinition(got, want)
}

func sameDefinition(a, b orbit.ScheduleDefinition) bool {
	// 原始JSON可有空白差异；模板和结果保持精确整数，不经float64。
	pa, ok := decodePayload(a.Task.Payload)
	pb, ok2 := decodePayload(b.Task.Payload)
	a.Task.Payload = nil
	b.Task.Payload = nil
	return ok && ok2 && pa == pb && reflect.DeepEqual(a, b)
}

func pauseSchedule(ctx context.Context, cfg config, api scheduleAPI, runID string, p *planReport) string {
	for attempt := 0; attempt < 3; attempt++ {
		call, stop := context.WithTimeout(ctx, cfg.Timeout)
		current, err := api.GetSchedule(call, p.ID)
		stop()
		if err != nil {
			return "PAUSE_READ_" + classify(err)
		}
		if current.ID != p.ID || current.Name != operationKey(runID, p.Index, "create") || current.QueueID != cfg.QueueID {
			return "PAUSE_IDENTITY_MISMATCH"
		}
		if current.Paused {
			p.PauseConfirmed = true
			return ""
		}
		p.PauseKey = operationKey(runID, p.Index, fmt.Sprintf("pause-%d", attempt))
		p.PauseRevision = current.Revision
		call, stop = context.WithTimeout(ctx, cfg.Timeout)
		res, err := api.PauseSchedule(call, p.ID, p.PauseKey, orbit.ControlScheduleRequest{ExpectedRevision: p.PauseRevision, Reason: "stop bounded schedule verification"})
		stop()
		if err == nil {
			if res.Schedule.ID != p.ID || !res.Schedule.Paused {
				return "PAUSE_RECEIPT_MISMATCH"
			}
			p.PauseConfirmed = true
			return ""
		}
		var apiErr *orbit.APIError
		if !errors.As(err, &apiErr) || apiErr.Code != "STATE_CONFLICT" || apiErr.StatusCode != 409 {
			return "PAUSE_" + classify(err)
		}
		timer := time.NewTimer(25 * time.Millisecond)
		select {
		case <-ctx.Done():
			timer.Stop()
			return "PAUSE_" + classify(ctx.Err())
		case <-timer.C:
		}
	}
	return "PAUSE_CONFLICT_LIMIT"
}

func validCalendarVersion(v string) bool {
	if len(v) == 0 || len(v) > 64 {
		return false
	}
	for _, c := range v {
		if !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || c == '.' || c == '-' || c == '_') {
			return false
		}
	}
	return true
}
func sameVersion(s orbit.Schedule, p planReport) bool {
	return s.Current.Definition.Rule.CalendarVersion == p.CalendarVersion && s.Current.EffectiveAt.Equal(p.EffectiveAt)
}