package main

import (
	"context"
	"encoding/json"
	"fmt"
	orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
	"sync"
	"testing"
	"time"
)

func testID(n int) string { return fmt.Sprintf("%032x", n) }
func testConfig() config {
	return config{QueueID: testID(1), Requests: 1, Concurrency: 1, Workers: 1, Capacity: 2, Timeout: time.Second, RunTimeout: 90 * time.Second, PollInterval: 100 * time.Millisecond}
}

type proofAPI struct {
	workflowAPI
	instance orbit.WorkflowInstance
	steps    orbit.WorkflowSteps
	tasks    map[string]orbit.Task
	attempts map[string]orbit.AttemptPage
	sources  map[string]orbit.TaskSource
}

func (a *proofAPI) GetWorkflow(context.Context, string) (orbit.WorkflowInstance, error) {
	return a.instance, nil
}
func (a *proofAPI) GetWorkflowSteps(context.Context, string) (orbit.WorkflowSteps, error) {
	return a.steps, nil
}
func (a *proofAPI) GetTask(_ context.Context, id string) (orbit.Task, error) { return a.tasks[id], nil }
func (a *proofAPI) ListTaskAttempts(_ context.Context, id string, _ orbit.HistoryPageRequest) (orbit.AttemptPage, error) {
	return a.attempts[id], nil
}
func (a *proofAPI) GetTaskSource(_ context.Context, id string) (orbit.TaskSource, error) {
	return a.sources[id], nil
}
func (a *proofAPI) StartWorkflow(context.Context, string, orbit.StartWorkflowRequest) (orbit.WorkflowStartResult, error) {
	return orbit.WorkflowStartResult{Instance: a.instance}, nil
}
func proofFixture() (*proofAPI, workflowAdmission) {
	c := testConfig()
	start := time.Date(2026, 9, 20, 0, 0, 0, 0, time.UTC)
	end := start.Add(3 * time.Second)
	raw := json.RawMessage(`{"run_id":"run","sequence":0}`)
	instance := orbit.WorkflowInstance{ID: testID(2), NamespaceID: testID(3), DefinitionID: testID(4), DefinitionVersionID: testID(5), Version: "v1", State: "SUCCEEDED", Revision: 4, Input: raw, Output: json.RawMessage(`{"left":{"run_id":"run","sequence":0},"right":{"run_id":"run","sequence":0}}`), BusinessKey: businessKey("run", 0), DeadlineAt: start.Add(c.RunTimeout), CreatedAt: start, CompletedAt: &end}
	api := &proofAPI{instance: instance, steps: orbit.WorkflowSteps{InstanceID: instance.ID, Revision: instance.Revision}, tasks: map[string]orbit.Task{}, attempts: map[string]orbit.AttemptPage{}, sources: map[string]orbit.TaskSource{}}
	for i, key := range []string{"pair", "left", "right"} {
		activation := testID(10 + i)
		step := orbit.WorkflowStepStatus{ID: testID(20 + i), Key: key, Kind: "TASK", State: "SUCCEEDED", Revision: 3, ActivationID: &activation, ActivationSequence: 1}
		if i == 0 {
			step.Kind = "PARALLEL"
			api.steps.Steps = append(api.steps.Steps, step)
			continue
		}
		taskID := testID(30 + i)
		step.TaskID = &taskID
		api.steps.Steps = append(api.steps.Steps, step)
		authorized, completed := start.Add(time.Second), start.Add(2*time.Second)
		task := orbit.Task{ID: taskID, NamespaceID: instance.NamespaceID, QueueID: c.QueueID, Handler: handlerName("run"), HandlerVersion: "v1", PayloadVersion: "v1", State: orbit.TaskSucceeded, Payload: raw, Result: raw, RecoverySafety: orbit.RecoveryReconcileOnly, Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: uint64(c.RunTimeout.Milliseconds())}, AttemptCount: 1, CreatedAt: start, AvailableAt: start, DeadlineAt: instance.DeadlineAt, CompletedAt: &completed}
		api.tasks[taskID] = task
		api.attempts[taskID] = orbit.AttemptPage{Items: []orbit.AttemptSummary{{ID: testID(40 + i), TaskID: taskID, State: orbit.AttemptSucceeded, Epoch: 1, CreatedAt: start, AuthorizedAt: &authorized, EndedAt: &completed}}}
		api.sources[taskID] = orbit.TaskSource{TaskID: taskID, Kind: "WORKFLOW", Workflow: &orbit.TaskWorkflowSource{WorkflowID: instance.ID, DefinitionID: instance.DefinitionID, DefinitionVersionID: instance.DefinitionVersionID, Version: "v1", StepID: step.ID, StepKey: key, StepKind: "TASK", ActivationID: activation, ActivationSequence: 1}}
	}
	return api, workflowAdmission{sequence: 0, instance: instance}
}
func TestWorkflowProofComplete(t *testing.T) {
	api, admission := proofFixture()
	cfg := testConfig()
	duration, proofs, code := verifyWorkflow(context.Background(), cfg, api, "run", admission.instance.DefinitionID, admission)
	if code != "" || duration != 3*time.Second || len(proofs) != 2 || proofs[0].taskID == proofs[1].taskID {
		t.Fatalf("valid proof rejected %s %#v", code, proofs)
	}
	log := newInvocationLog(1)
	handler := log.handler(context.Background(), cfg, "run")
	var wg sync.WaitGroup
	for _, p := range proofs {
		wg.Add(1)
		go func(p verifiedTask) {
			defer wg.Done()
			_, err := handler.Handle(context.Background(), orbit.WorkerExecution{Task: api.tasks[p.taskID], AttemptID: p.attemptID, Epoch: 1})
			if err != nil {
				t.Error(err)
			}
		}(p)
	}
	wg.Wait()
	r := newReport(cfg, "run")
	log.summarize(&r, proofs)
	if r.HandlerCalls != 2 || r.DuplicateInvocations != 0 || r.UnconfirmedInvocations != 0 || len(r.finish().Errors) != 0 {
		t.Fatalf("call proof mismatch %#v", r)
	}
}
func TestWorkflowProofRejectsWrongIdentitiesAndHistory(t *testing.T) {
	for name, change := range map[string]func(*proofAPI){
		"version substitution": func(a *proofAPI) { a.instance.DefinitionVersionID = testID(80) },
		"changed input":        func(a *proofAPI) { a.instance.Input = json.RawMessage(`{"run_id":"run"}`) },
		"wrong steps owner":    func(a *proofAPI) { a.steps.InstanceID = testID(80) },
		"missing join":         func(a *proofAPI) { a.steps.Steps = a.steps.Steps[1:] },
		"duplicate branch":     func(a *proofAPI) { a.steps.Steps[2] = a.steps.Steps[1] },
		"wrong branch name":    func(a *proofAPI) { a.steps.Steps[2].Key = "third" },
		"reused task":          func(a *proofAPI) { a.steps.Steps[2].TaskID = a.steps.Steps[1].TaskID },
		"reused activation":    func(a *proofAPI) { a.steps.Steps[2].ActivationID = a.steps.Steps[1].ActivationID },
		"wrong task policy":    func(a *proofAPI) { v := a.tasks[testID(31)]; v.Retry.MaxAttempts = 2; a.tasks[v.ID] = v },
		"wrong task queue":     func(a *proofAPI) { v := a.tasks[testID(31)]; v.QueueID = testID(80); a.tasks[v.ID] = v },
		"missing result sequence": func(a *proofAPI) {
			v := a.tasks[testID(31)]
			v.Result = json.RawMessage(`{"run_id":"run"}`)
			a.tasks[v.ID] = v
		},
		"uncertain task": func(a *proofAPI) { v := a.tasks[testID(31)]; v.EffectUncertain = true; a.tasks[v.ID] = v },
		"wrong task deadline": func(a *proofAPI) {
			v := a.tasks[testID(31)]
			v.DeadlineAt = v.DeadlineAt.Add(time.Second)
			a.tasks[v.ID] = v
		},
		"extra attempt": func(a *proofAPI) {
			v := a.attempts[testID(31)]
			v.Items = append(v.Items, v.Items[0])
			a.attempts[testID(31)] = v
		},
		"missing authorization": func(a *proofAPI) { a.attempts[testID(31)].Items[0].AuthorizedAt = nil },
		"wrong epoch":           func(a *proofAPI) { a.attempts[testID(31)].Items[0].Epoch = 2 },
		"wrong attempt time": func(a *proofAPI) {
			v := a.instance.CreatedAt.Add(-time.Second)
			a.attempts[testID(31)].Items[0].EndedAt = &v
		},
		"wrong source task":   func(a *proofAPI) { v := a.sources[testID(31)]; v.TaskID = testID(80); a.sources[testID(31)] = v },
		"wrong source branch": func(a *proofAPI) { a.sources[testID(31)].Workflow.StepKey = "right" },
		"wrong source kind":   func(a *proofAPI) { a.sources[testID(31)].Workflow.StepKind = "POLL" },
		"wrong output zero": func(a *proofAPI) {
			a.instance.Output = json.RawMessage(`{"left":{"run_id":"run"},"right":{"run_id":"run"}}`)
		},
	} {
		t.Run(name, func(t *testing.T) {
			a, admission := proofFixture()
			change(a)
			_, proofs, code := verifyWorkflow(context.Background(), testConfig(), a, "run", admission.instance.DefinitionID, admission)
			if code == "" || len(proofs) != 0 {
				t.Fatalf("bad evidence counted %s %#v", code, proofs)
			}
		})
	}
}
func TestStartBindsPublishedVersion(t *testing.T) {
	a, _ := proofFixture()
	c := testConfig()
	results := submitWorkflows(context.Background(), c, a, "run", a.instance.DefinitionID, a.instance.DefinitionVersionID, a.instance.NamespaceID)
	if len(results) != 1 || results[0].code != "" {
		t.Fatalf("valid start rejected %#v", results)
	}
	results = submitWorkflows(context.Background(), c, a, "run", a.instance.DefinitionID, testID(80), a.instance.NamespaceID)
	if len(results) != 1 || results[0].code != "INVALID_RECEIPT" {
		t.Fatal("unpublished version accepted")
	}
}
func TestDuplicateInvocationCannotCountAsSecondBranch(t *testing.T) {
	a, _ := proofFixture()
	c := testConfig()
	log := newInvocationLog(1)
	h := log.handler(context.Background(), c, "run")
	proof := verifiedTask{sequence: 0, taskID: testID(31), attemptID: testID(41)}
	for i := 0; i < 2; i++ {
		_, _ = h.Handle(context.Background(), orbit.WorkerExecution{Task: a.tasks[proof.taskID], AttemptID: proof.attemptID, Epoch: 1})
	}
	r := newReport(c, "run")
	log.summarize(&r, []verifiedTask{proof, {sequence: 0, taskID: testID(32), attemptID: testID(42)}})
	if r.HandlerCalls != 2 || r.DuplicateInvocations != 1 || r.Errors["LOCAL_PROOF_MISMATCH"] != 1 {
		t.Fatalf("duplicate filled missing branch %#v", r)
	}
	r = newReport(c, "run")
	log.summarize(&r, nil)
	if r.UnconfirmedInvocations != 2 {
		t.Fatal("unverified calls disappeared")
	}
}
func TestStrictOutputKeepsSequenceZeroRequired(t *testing.T) {
	for _, raw := range []string{
		`{"left":{"run_id":"run"},"right":{"run_id":"run","sequence":0}}`,
		`{"left":{"run_id":"run","sequence":null},"right":{"run_id":"run","sequence":0}}`,
		`{"left":{"run_id":"run","sequence":0,"sequence":0},"right":{"run_id":"run","sequence":0}}`,
		`{"left":{"run_id":"run","Sequence":0},"right":{"run_id":"run","sequence":0}}`,
		`{"left":{"run_id":"run","sequence":0},"right":{"run_id":"run","sequence":0},"right":{"run_id":"run","sequence":0}}`,
	} {
		if matchesOutput(json.RawMessage(raw), "run", 0) {
			t.Fatalf("invalid zero payload accepted %s", raw)
		}
	}
}
func TestSpecificationRetainsConfiguredTimeout(t *testing.T) {
	c := testConfig()
	spec := workflowSpec(c.QueueID, handlerName("run"), c.RunTimeout)
	if spec.TimeoutMS != 90000 || spec.Body.Steps[0].Parallel.Branches[0].Body.Steps[0].Task.TimeoutMS != 90000 {
		t.Fatal("run timeout silently capped")
	}
	var stored orbit.WorkflowSpecification
	raw, _ := json.Marshal(spec)
	if err := json.Unmarshal(raw, &stored); err != nil {
		t.Fatal(err)
	}
	if !sameSpecification(stored, spec) {
		t.Fatal("serialized specification rejected")
	}
	stored.TimeoutMS = 30000
	if sameSpecification(stored, spec) {
		t.Fatal("changed published spec accepted")
	}
}