package main

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

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

// workflowPayload 是两个并行分支共享的固定输入；不含凭据或外部地址。
type workflowPayload struct {
	RunID    string `json:"run_id"`   // RunID 隔离本次测量，不作为业务幂等保证。
	Sequence int    `json:"sequence"` // Sequence 从0开始的实例序号，用于发现错配和重复调用。
}

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

type workflowAPI interface {
	GetQueue(context.Context, string) (orbit.Queue, error)
	CreateWorkflowDefinition(context.Context, string, orbit.CreateWorkflowDefinitionRequest) (orbit.WorkflowDefinitionResult, error)
	PublishWorkflowVersion(context.Context, string, string, orbit.PublishWorkflowVersionRequest) (orbit.WorkflowVersionResult, error)
	StartWorkflow(context.Context, string, orbit.StartWorkflowRequest) (orbit.WorkflowStartResult, error)
	GetWorkflow(context.Context, string) (orbit.WorkflowInstance, error)
	GetWorkflowSteps(context.Context, string) (orbit.WorkflowSteps, error)
	GetTask(context.Context, string) (orbit.Task, error)
	ListTaskAttempts(context.Context, string, orbit.HistoryPageRequest) (orbit.AttemptPage, error)
	GetTaskSource(context.Context, string) (orbit.TaskSource, error)
}

type workflowAdmission struct {
	sequence int
	instance orbit.WorkflowInstance
	duration time.Duration
	code     string
}

type workflowVerification struct {
	sequence int
	duration time.Duration
	code     string
	tasks    []verifiedTask
}

func workflowSpec(queue, handler string, timeout time.Duration) orbit.WorkflowSpecification {
	typ := orbit.WorkflowType{Kind: "OBJECT", Fields: map[string]orbit.WorkflowType{"run_id": {Kind: "STRING"}, "sequence": {Kind: "NUMBER"}}}
	task := orbit.WorkflowTask{QueueID: queue, Handler: handler, HandlerVersion: "v1", PayloadVersion: "v1", InputType: typ, Input: orbit.WorkflowRef("input"), ResultType: typ, TimeoutMS: uint64(timeout.Milliseconds()), Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: uint64(timeout.Milliseconds())}, RecoverySafety: orbit.RecoveryReconcileOnly}
	parallel := orbit.WorkflowParallel{Branches: []orbit.WorkflowParallelBranch{
		{Name: "left", Body: orbit.WorkflowBlock{Steps: []orbit.WorkflowStep{{Key: "left", Kind: "TASK", Task: &task}}, OutputType: typ, Output: orbit.WorkflowRef("left")}},
		{Name: "right", Body: orbit.WorkflowBlock{Steps: []orbit.WorkflowStep{{Key: "right", Kind: "TASK", Task: &task}}, OutputType: typ, Output: orbit.WorkflowRef("right")}},
	}}
	output := orbit.WorkflowType{Kind: "OBJECT", Fields: map[string]orbit.WorkflowType{"left": typ, "right": typ}}
	return orbit.WorkflowSpecification{SchemaVersion: 1, InputType: typ, TimeoutMS: uint64(timeout.Milliseconds()), FailurePolicy: "REVIEW", Body: orbit.WorkflowBlock{Steps: []orbit.WorkflowStep{{Key: "pair", Kind: "PARALLEL", Parallel: &parallel}}, OutputType: output, Output: orbit.WorkflowRef("pair")}}
}

func measure(parent context.Context, cfg config, api *orbit.Client, runID string) workflowReport {
	return measureClient(parent, cfg, api, runID)
}

func measureClient(parent context.Context, cfg config, api *orbit.Client, runID string) workflowReport {
	r := newReport(cfg, runID)
	ctx, cancel := context.WithTimeout(parent, cfg.RunTimeout)
	defer cancel()
	requestCtx, stop := context.WithTimeout(ctx, cfg.Timeout)
	queue, err := api.GetQueue(requestCtx, cfg.QueueID)
	stop()
	if err != nil {
		r.Errors["PREFLIGHT_"+classify(err)]++
		return r.finish()
	}
	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.finish()
	}
	handler := handlerName(runID)
	spec := workflowSpec(cfg.QueueID, handler, cfg.RunTimeout)
	requestCtx, stop = context.WithTimeout(ctx, cfg.Timeout)
	definition, err := api.CreateWorkflowDefinition(requestCtx, definitionKey(runID), orbit.CreateWorkflowDefinitionRequest{Name: definitionKey(runID), Description: "bounded parallel workflow verification"})
	stop()
	if err != nil {
		r.Errors["CREATE_"+classify(err)]++
		return r.finish()
	}
	r.DefinitionID = definition.Definition.ID
	if !validID(r.DefinitionID) || definition.Deduplicated || definition.Definition.NamespaceID != queue.NamespaceID || definition.Definition.Name != definitionKey(runID) {
		r.Errors["INVALID_DEFINITION_RECEIPT"]++
		return r.finish()
	}
	requestCtx, stop = context.WithTimeout(ctx, cfg.Timeout)
	version, err := api.PublishWorkflowVersion(requestCtx, r.DefinitionID, runID, orbit.PublishWorkflowVersionRequest{Version: "v1", Specification: spec, Reason: "bounded parallel workflow verification"})
	stop()
	if err != nil {
		r.Errors["PUBLISH_"+classify(err)]++
		return r.finish()
	}
	r.VersionID = version.Version.ID
	if !validID(r.VersionID) || version.Version.Version != "v1" || version.Version.DefinitionID != r.DefinitionID || version.Deduplicated || version.Version.NamespaceID != queue.NamespaceID || !sameSpecification(version.Version.Specification, spec) {
		r.Errors["INVALID_VERSION_RECEIPT"]++
		return r.finish()
	}

	log := newInvocationLog(cfg.Requests)
	workers := make([]*orbit.Worker, 0, cfg.Workers)
	for i := 0; i < cfg.Workers; i++ {
		worker, 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 r.finish()
		}
		workers = append(workers, worker)
	}
	workerCtx, stopWorkers := context.WithCancel(ctx)
	done := make(chan error, len(workers))
	for _, worker := range workers {
		go func(w *orbit.Worker) {
			e := w.Run(workerCtx)
			if workerCtx.Err() == nil {
				cancel()
			}
			done <- e
		}(worker)
	}

	admissions := submitWorkflows(ctx, cfg, api, runID, r.DefinitionID, r.VersionID, queue.NamespaceID)
	r.Attempted = len(admissions)
	r.NotSubmitted = cfg.Requests - r.Attempted
	var accepted []workflowAdmission
	var admissionDurations []time.Duration
	seenInstances := map[string]bool{}
	for _, a := range admissions {
		admissionDurations = append(admissionDurations, a.duration)
		if a.code == "" && seenInstances[a.instance.ID] {
			a.code = "DUPLICATE_INSTANCE"
		}
		if a.code == "" {
			seenInstances[a.instance.ID] = true
		}
		r.Instances = append(r.Instances, instanceReport{Sequence: a.sequence, Key: requestKey(runID, a.sequence), ID: a.instance.ID, Code: a.code})
		if a.code != "" {
			r.AdmissionErrors++
			r.Errors["ADMISSION_"+a.code]++
			continue
		}
		accepted = append(accepted, a)
	}
	r.Accepted = len(accepted)
	r.Admission = distributionOf(admissionDurations)
	r.AdmissionFinishedAt = time.Now().UTC()
	verified := verifyWorkflows(ctx, cfg, api, runID, r.DefinitionID, accepted)
	var completions []time.Duration
	var proven []verifiedTask
	for _, v := range verified {
		if v.code != "" {
			r.Errors[v.code]++
			continue
		}
		r.VerifiedSuccess++
		proven = append(proven, v.tasks...)
		for i := range r.Instances {
			if r.Instances[i].Sequence == v.sequence {
				r.Instances[i].Verified = true
				break
			}
		}
		completions = append(completions, v.duration)
	}
	if ctx.Err() != nil {
		r.Errors["RUN_"+classify(ctx.Err())]++
	}
	stopWorkers()
	for range workers {
		if e := <-done; e != nil {
			r.Errors["WORKER_STOP"]++
		}
	}
	for _, worker := range workers {
		stats := worker.Stats()
		r.WorkerStarted += stats.Started
		r.Reconnects += stats.Reconnects
		r.Renewals += stats.RenewalsConfirmed
	}
	log.summarize(&r, proven)
	r.Completion = distributionOf(completions)
	return r.finish()
}

func submitWorkflows(ctx context.Context, cfg config, api workflowAPI, runID, definitionID, versionID, namespace string) []workflowAdmission {
	jobs := make(chan int)
	out := make(chan workflowAdmission, 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(workflowPayload{RunID: runID, Sequence: sequence})
				started := time.Now()
				call, stop := context.WithTimeout(ctx, cfg.Timeout)
				result, err := api.StartWorkflow(call, requestKey(runID, sequence), orbit.StartWorkflowRequest{DefinitionID: definitionID, Version: "v1", Input: payload, BusinessKey: businessKey(runID, sequence)})
				stop()
				a := workflowAdmission{sequence: sequence, instance: result.Instance, duration: time.Since(started)}
				if err != nil {
					a.code = classify(err)
				} else if result.Deduplicated || !matchesInstance(result.Instance, cfg, runID, sequence, definitionID, versionID, namespace) {
					a.code = "INVALID_RECEIPT"
				}
				out <- a
			}
		}()
	}
feed:
	for sequence := 0; sequence < cfg.Requests; sequence++ {
		select {
		case <-ctx.Done():
			break feed
		case jobs <- sequence:
		}
	}
	close(jobs)
	wg.Wait()
	close(out)
	results := make([]workflowAdmission, 0, cfg.Requests)
	for result := range out {
		results = append(results, result)
	}
	return results
}

func verifyWorkflows(ctx context.Context, cfg config, api workflowAPI, runID, definitionID string, admissions []workflowAdmission) []workflowVerification {
	jobs := make(chan workflowAdmission, len(admissions))
	for _, admission := range admissions {
		jobs <- admission
	}
	close(jobs)
	out := make(chan workflowVerification, len(admissions))
	var wg sync.WaitGroup
	for i := 0; i < cfg.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for admission := range jobs {
				duration, tasks, code := verifyWorkflow(ctx, cfg, api, runID, definitionID, admission)
				out <- workflowVerification{sequence: admission.sequence, duration: duration, tasks: tasks, code: code}
			}
		}()
	}
	wg.Wait()
	close(out)
	results := make([]workflowVerification, 0, len(admissions))
	for result := range out {
		results = append(results, result)
	}
	return results
}

func verifyWorkflow(ctx context.Context, cfg config, api workflowAPI, runID, definitionID string, admission workflowAdmission) (time.Duration, []verifiedTask, string) {
	for {
		if ctx.Err() != nil {
			return 0, nil, "VERIFY_" + classify(ctx.Err())
		}
		call, stop := context.WithTimeout(ctx, cfg.Timeout)
		instance, err := api.GetWorkflow(call, admission.instance.ID)
		stop()
		if err != nil {
			return 0, nil, "READ_" + classify(err)
		}
		if instance.ID != admission.instance.ID || !matchesInstance(instance, cfg, runID, admission.sequence, definitionID, admission.instance.DefinitionVersionID, admission.instance.NamespaceID) || !instance.CreatedAt.Equal(admission.instance.CreatedAt) || !instance.DeadlineAt.Equal(admission.instance.DeadlineAt) {
			return 0, nil, "WORKFLOW_MISMATCH"
		}
		switch instance.State {
		case "SUCCEEDED":
			if instance.CompletedAt == nil || instance.CreatedAt.IsZero() || instance.CompletedAt.Before(instance.CreatedAt) {
				return 0, nil, "COMPLETION_MISMATCH"
			}
			call, stop = context.WithTimeout(ctx, cfg.Timeout)
			steps, e := api.GetWorkflowSteps(call, instance.ID)
			stop()
			if e != nil {
				return 0, nil, "STEPS_" + classify(e)
			}
			tasks, code := verifySteps(ctx, cfg, api, instance, runID, admission.sequence, steps)
			if code != "" {
				return 0, nil, code
			}
			if !matchesOutput(instance.Output, runID, admission.sequence) {
				return 0, nil, "OUTPUT_MISMATCH"
			}
			return instance.CompletedAt.Sub(instance.CreatedAt), tasks, ""
		case "NEEDS_REVIEW", "FAILED", "CANCELLED":
			return 0, nil, "WORKFLOW_" + instance.State
		case "RUNNING", "WAITING", "PAUSED", "PENDING":
		default:
			return 0, nil, "UNKNOWN_WORKFLOW_STATE"
		}
		timer := time.NewTimer(cfg.PollInterval)
		select {
		case <-ctx.Done():
			timer.Stop()
			return 0, nil, "VERIFY_" + classify(ctx.Err())
		case <-timer.C:
		}
	}
}

func sameSpecification(a, b orbit.WorkflowSpecification) bool {
	left, e := json.Marshal(a)
	right, e2 := json.Marshal(b)
	return e == nil && e2 == nil && bytes.Equal(left, right)
}
func matchesInstance(i orbit.WorkflowInstance, c config, runID string, sequence int, definitionID, versionID, namespace string) bool {
	p, ok := decodePayload(i.Input)
	return validID(i.ID) && i.NamespaceID == namespace && i.DefinitionID == definitionID && i.DefinitionVersionID == versionID && i.Version == "v1" && i.BusinessKey == businessKey(runID, sequence) && ok && p == (workflowPayload{RunID: runID, Sequence: sequence}) && !i.CreatedAt.IsZero() && i.DeadlineAt.Equal(i.CreatedAt.Add(time.Duration(c.RunTimeout.Milliseconds())*time.Millisecond)) && !i.Paused && i.CancelRequestedAt == nil && i.StopKind == nil && i.StopRequestedAt == nil
}