Agent-System/internal/projectinstance/store_mem.go

package projectinstance

import (
	"context"
	"fmt"
	"sync"
	"time"

	"chat-agent-system/internal/workflowdef"
)

// =============================================================================
// 内存 Store（单元测试 / 无 PG 迁移）
// =============================================================================

type memStageRecord struct {
	id       int64
	stageKey string
	position int
	enabled  bool
	status   string
	tasks    []memDocTaskRecord
}

type memDocTaskRecord struct {
	id      int64
	taskKey string
	enabled bool
	status  string
}

type memStore struct {
	mu          sync.Mutex
	meta        map[int64]projectMeta
	plan        map[int64]*WorkflowInstance
	live        map[int64]*WorkflowInstance
	tokens      map[string]activationTokenRec
	stages      map[int64][]memStageRecord
	nextStageID int64
	nextTaskID  int64
}

type projectMeta struct {
	lifecycle    string
	activeKey    string
	selectionVer int
	planRevision int
}

type activationTokenRec struct {
	projectID      int64
	planHash       string
	expectedVer    int
	snapshot       WorkflowPlan
	expiresAt      time.Time
	used           bool
	policyRevision *int64
}

func newMemStore() *memStore {
	return &memStore{
		meta:        make(map[int64]projectMeta),
		plan:        make(map[int64]*WorkflowInstance),
		live:        make(map[int64]*WorkflowInstance),
		tokens:      make(map[string]activationTokenRec),
		stages:      make(map[int64][]memStageRecord),
		nextStageID: 1,
		nextTaskID:  1,
	}
}

func (m *memStore) rebuildStagesLocked(inst *WorkflowInstance, activate bool) {
	records := make([]memStageRecord, 0, len(inst.ConfigSnapshot.Stages))
	firstActive := true
	for i, st := range inst.ConfigSnapshot.Stages {
		rec := memStageRecord{
			id: m.nextStageID, stageKey: st.StageKey, position: i, enabled: st.Enabled,
		}
		m.nextStageID++
		if !st.Enabled {
			rec.status = StageSkipped
		} else if activate && firstActive {
			rec.status = StageActive
			firstActive = false
		} else {
			rec.status = StagePending
		}
		for _, dt := range st.DocTasks {
			task := memDocTaskRecord{id: m.nextTaskID, taskKey: dt.TaskKey, enabled: dt.Enabled, status: "pending"}
			m.nextTaskID++
			if !dt.Enabled {
				task.status = "skipped_disabled"
			}
			rec.tasks = append(rec.tasks, task)
		}
		records = append(records, rec)
	}
	m.stages[inst.ID] = records
}

func (m *memStore) instanceDetailLocked(inst *WorkflowInstance) *InstanceDetail {
	detail := &InstanceDetail{Instance: *inst, Stages: make([]StageInstance, 0)}
	for _, rec := range m.stages[inst.ID] {
		st := StageInstance{
			ID: rec.id, StageKey: rec.stageKey, Position: rec.position, Enabled: rec.enabled, Status: rec.status,
		}
		if st.Status == StageActive {
			detail.ActiveStageKey = st.StageKey
		}
		for _, t := range rec.tasks {
			st.DocTasks = append(st.DocTasks, DocTaskInstance{
				ID: t.id, TaskKey: t.taskKey, Enabled: t.enabled, Status: t.status, Approved: t.status == "approved",
			})
		}
		detail.Stages = append(detail.Stages, st)
	}
	return detail
}

func (m *memStore) ensureMeta(projectID int64) projectMeta {
	if pm, ok := m.meta[projectID]; ok {
		return pm
	}
	pm := projectMeta{lifecycle: LifecyclePlanning, selectionVer: 1}
	m.meta[projectID] = pm
	return pm
}

func (m *memStore) GetProjectMeta(_ context.Context, projectID int64) (string, string, int, int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	pm, ok := m.meta[projectID]
	if !ok {
		return "", "", 0, 0, ErrProjectNotFound
	}
	return pm.lifecycle, pm.activeKey, pm.selectionVer, pm.planRevision, nil
}

func (m *memStore) UpdateProjectAfterActivate(_ context.Context, projectID int64, workflowKey string, expectedVersion int, _ int64) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	pm := m.ensureMeta(projectID)
	if pm.selectionVer != expectedVersion {
		return ErrConcurrentConflict
	}
	pm.lifecycle = LifecycleActive
	pm.activeKey = workflowKey
	pm.selectionVer++
	m.meta[projectID] = pm
	return nil
}

func (m *memStore) BumpPlanRevision(_ context.Context, projectID int64, expectedPlanRevision int) (int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	pm := m.ensureMeta(projectID)
	if pm.planRevision != expectedPlanRevision {
		return 0, ErrPlanRevisionConflict
	}
	pm.planRevision++
	m.meta[projectID] = pm
	return pm.planRevision, nil
}

func (m *memStore) GetPlanningInstance(_ context.Context, projectID int64) (*WorkflowInstance, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	inst := m.plan[projectID]
	if inst == nil {
		return nil, nil
	}
	copyInst := *inst
	return &copyInst, nil
}

func (m *memStore) GetLiveInstance(_ context.Context, projectID int64) (*WorkflowInstance, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	inst := m.live[projectID]
	if inst == nil || !IsLiveInstanceStatus(inst.Status) {
		return nil, nil
	}
	copyInst := *inst
	return &copyInst, nil
}

func (m *memStore) GetGoverningInstanceForAccess(_ context.Context, projectID int64) (*WorkflowInstance, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	inst := m.live[projectID]
	if inst == nil {
		return nil, nil
	}
	if IsLiveInstanceStatus(inst.Status) || IsClosedInstanceStatus(inst.Status) {
		copyInst := *inst
		return &copyInst, nil
	}
	// planning 等非治理态：视为从未激活
	return nil, nil
}

func (m *memStore) HasLiveInstance(_ context.Context, projectID int64) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	inst := m.live[projectID]
	return inst != nil && IsLiveInstanceStatus(inst.Status), nil
}

func (m *memStore) UpsertPlanningInstance(_ context.Context, inst *WorkflowInstance) (*WorkflowInstance, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if inst.ID == 0 {
		inst.ID = int64(len(m.plan) + 1)
	}
	inst.Status = InstancePlanning
	now := time.Now()
	if inst.CreatedAt.IsZero() {
		inst.CreatedAt = now
	}
	inst.UpdatedAt = now
	copyInst := *inst
	m.plan[inst.ProjectID] = &copyInst
	m.rebuildStagesLocked(&copyInst, false)
	return &copyInst, nil
}

// SavePlanWithRevision 内存实现：bump 与 upsert 在同一互斥锁内，天然原子。
func (m *memStore) SavePlanWithRevision(ctx context.Context, inst *WorkflowInstance, expectedPlanRevision int) (*WorkflowInstance, int, error) {
	newRev, err := m.BumpPlanRevision(ctx, inst.ProjectID, expectedPlanRevision)
	if err != nil {
		return nil, 0, err
	}
	saved, err := m.UpsertPlanningInstance(ctx, inst)
	if err != nil {
		return nil, 0, err
	}
	return saved, newRev, nil
}

func (m *memStore) ActivateInstanceWithToken(ctx context.Context, token string, projectID int64, inst *WorkflowInstance, version *workflowdef.Version, expectedVersion int, activatedBy int64, planHash string) (*WorkflowInstance, error) {
	usesAuthorOnly := planUsesAuthorOnly(inst.ConfigSnapshot)
	var currentPolicyRev int64
	if usesAuthorOnly {
		snap, lockErr, ok := runActivateAuthorOnlyLock(ctx, nil)
		if !ok {
			return nil, ErrAuthorOnlyPolicyUnavailable
		}
		if lockErr != nil {
			return nil, fmt.Errorf("%w: %v", ErrAuthorOnlyPolicyUnavailable, lockErr)
		}
		if !snap.Allowed {
			return nil, ErrAuthorOnlyPolicyDisabled
		}
		currentPolicyRev = snap.Revision
	}
	if err := runActivateMemberRecheck(ctx); err != nil {
		return nil, err
	}
	tokenPolicyRev, err := m.consumeActivationTokenLocked(ctx, token, projectID, planHash)
	if err != nil {
		return nil, err
	}
	if usesAuthorOnly {
		if tokenPolicyRev == nil || *tokenPolicyRev != currentPolicyRev {
			return nil, ErrAuthorOnlyPolicyChanged
		}
	}
	return m.ActivateInstance(ctx, projectID, inst, version, expectedVersion, activatedBy)
}

func (m *memStore) CountActiveTaskRuns(_ context.Context, _ int64) (int, error) {
	return 0, nil
}

func (m *memStore) ActivateInstance(_ context.Context, projectID int64, inst *WorkflowInstance, version *workflowdef.Version, expectedVersion int, activatedBy int64) (*WorkflowInstance, error) {
	workflowKey := ""
	if version != nil {
		workflowKey = version.WorkflowKey
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	pm := m.ensureMeta(projectID)
	if pm.selectionVer != expectedVersion {
		return nil, ErrConcurrentConflict
	}
	pm.lifecycle = LifecycleActive
	pm.activeKey = workflowKey
	pm.selectionVer++
	m.meta[projectID] = pm

	now := time.Now()
	activated := *inst
	activated.Status = InstanceActive
	activated.ActivatedBy = &activatedBy
	activated.ActivatedAt = &now
	m.live[projectID] = &activated
	m.rebuildStagesLocked(&activated, true)
	delete(m.plan, projectID)
	return &activated, nil
}

func (m *memStore) GetInstanceByID(_ context.Context, instanceID int64) (*InstanceDetail, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, inst := range m.live {
		if inst.ID == instanceID {
			return m.instanceDetailLocked(inst), nil
		}
	}
	for _, inst := range m.plan {
		if inst.ID == instanceID {
			return m.instanceDetailLocked(inst), nil
		}
	}
	return nil, ErrInstanceNotFound
}

func (m *memStore) AdvanceStage(_ context.Context, instanceID int64, stageKey string, expectedRevision int, _ int64) (*InstanceDetail, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var inst *WorkflowInstance
	for _, live := range m.live {
		if live.ID == instanceID {
			inst = live
			break
		}
	}
	if inst == nil {
		return nil, ErrInstanceNotFound
	}
	if inst.ConfigRevision != expectedRevision {
		return nil, ErrInstanceRevisionConflict
	}
	records := m.stages[instanceID]
	activeIdx := -1
	for i, rec := range records {
		if rec.stageKey == stageKey {
			if rec.status != StageActive {
				return nil, ErrStageNotActive
			}
			activeIdx = i
			break
		}
	}
	if activeIdx < 0 {
		return nil, ErrInvalidStage
	}
	records[activeIdx].status = StageCompleted
	anotherActive := false
	for i, rec := range records {
		if i != activeIdx && rec.enabled && rec.status == StageActive {
			anotherActive = true
			break
		}
	}
	nextIdx := -1
	if !anotherActive {
		for i, rec := range records {
			if rec.enabled && rec.status == StagePending {
				nextIdx = i
				break
			}
		}
	}
	inst.ConfigRevision++
	inst.UpdatedAt = time.Now()
	if anotherActive {
		inst.Status = InstanceActive
	} else if nextIdx < 0 {
		inst.Status = InstanceCompleted
		if pm := m.meta[inst.ProjectID]; pm.lifecycle != "" {
			pm.lifecycle = LifecyclePlanning
			pm.activeKey = ""
			m.meta[inst.ProjectID] = pm
		}
	} else {
		records[nextIdx].status = StageActive
	}
	m.stages[instanceID] = records
	copyInst := *inst
	m.live[inst.ProjectID] = &copyInst
	return m.instanceDetailLocked(&copyInst), nil
}

func (m *memStore) ReopenStage(_ context.Context, instanceID int64, stageKey string, docTaskKeys []string, _ int64) (*InstanceDetail, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var inst *WorkflowInstance
	for _, live := range m.live {
		if live.ID == instanceID {
			inst = live
			break
		}
	}
	if inst == nil {
		return nil, ErrInstanceNotFound
	}
	records := m.stages[instanceID]
	idx := -1
	for i, rec := range records {
		if rec.stageKey == stageKey {
			idx = i
			break
		}
	}
	if idx < 0 {
		return nil, ErrInvalidStage
	}
	if records[idx].status != StageCompleted {
		return nil, ErrStageNotCompleted
	}
	records[idx].status = StageActive
	wanted := map[string]bool{}
	for _, key := range docTaskKeys {
		wanted[key] = true
	}
	for j := range records[idx].tasks {
		if wanted[records[idx].tasks[j].taskKey] {
			records[idx].tasks[j].status = "in_progress"
		}
	}
	if inst.Status == InstanceCompleted {
		inst.Status = InstanceActive
	}
	inst.ConfigRevision++
	inst.UpdatedAt = time.Now()
	m.stages[instanceID] = records
	copyInst := *inst
	m.live[inst.ProjectID] = &copyInst
	return m.instanceDetailLocked(&copyInst), nil
}

func (m *memStore) MarkDocTaskApproved(_ context.Context, projectID int64, taskKey string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	live := m.live[projectID]
	if live == nil {
		return nil
	}
	records := m.stages[live.ID]
	for i := range records {
		for j := range records[i].tasks {
			if records[i].tasks[j].taskKey == taskKey && records[i].tasks[j].enabled {
				records[i].tasks[j].status = "approved"
			}
		}
	}
	m.stages[live.ID] = records
	return nil
}

func (m *memStore) UpdateLiveDocTaskDuties(_ context.Context, projectID int64, taskKey string, plan WorkflowPlan, _ int64) (*WorkflowInstance, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	live := m.live[projectID]
	if live == nil {
		return nil, ErrInstanceNotFound
	}
	if findDocTaskPlan(plan, taskKey) == nil {
		return nil, ErrDocTaskNotFound
	}
	planHash, err := computePlanHash(plan)
	if err != nil {
		return nil, err
	}
	copyInst := *live
	copyInst.ConfigSnapshot = plan
	copyInst.PlanHash = planHash
	copyInst.ConfigRevision++
	copyInst.UpdatedAt = time.Now()
	m.live[projectID] = &copyInst
	return &copyInst, nil
}

func (m *memStore) UpdateLivePlan(_ context.Context, projectID int64, plan WorkflowPlan, expectedConfigRevision int, _ int64) (*WorkflowInstance, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	live := m.live[projectID]
	if live == nil {
		return nil, ErrInstanceNotFound
	}
	if live.ConfigRevision != expectedConfigRevision {
		return nil, ErrInstanceRevisionConflict
	}
	planHash, err := computePlanHash(plan)
	if err != nil {
		return nil, err
	}
	copyInst := *live
	copyInst.ConfigSnapshot = plan
	copyInst.PlanHash = planHash
	copyInst.ConfigRevision++
	copyInst.UpdatedAt = time.Now()
	m.live[projectID] = &copyInst

	records := m.stages[copyInst.ID]
	for i := range records {
		for _, st := range plan.Stages {
			if records[i].stageKey == st.StageKey {
				records[i].enabled = st.Enabled
			}
		}
		for j := range records[i].tasks {
			dt := findDocTaskPlan(plan, records[i].tasks[j].taskKey)
			if dt == nil {
				continue
			}
			records[i].tasks[j].enabled = dt.Enabled
			if !dt.Enabled && records[i].tasks[j].status == "pending" {
				records[i].tasks[j].status = "skipped_disabled"
			}
			if dt.Enabled && records[i].tasks[j].status == "skipped_disabled" {
				records[i].tasks[j].status = "pending"
			}
		}
	}
	m.stages[copyInst.ID] = records
	return &copyInst, nil
}

func (m *memStore) TerminateInstance(_ context.Context, instanceID int64, projectID int64, terminatedBy int64, reason string) (*WorkflowInstance, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	inst := m.live[projectID]
	if inst == nil || inst.ID != instanceID {
		return nil, ErrInstanceNotFound
	}
	now := time.Now()
	inst.Status = InstanceTerminated
	copyInst := *inst
	m.live[projectID] = &copyInst
	pm := m.ensureMeta(projectID)
	pm.lifecycle = LifecyclePlanning
	pm.activeKey = ""
	m.meta[projectID] = pm
	_ = terminatedBy
	_ = reason
	_ = now
	return &copyInst, nil
}

func (m *memStore) SaveActivationToken(_ context.Context, token string, projectID int64, planHash string, expectedVersion int, snapshot WorkflowPlan, _ int64, expiresAt time.Time, policyRevision *int64) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	var revCopy *int64
	if policyRevision != nil {
		v := *policyRevision
		revCopy = &v
	}
	m.tokens[hashActivationToken(token)] = activationTokenRec{
		projectID: projectID, planHash: planHash, expectedVer: expectedVersion,
		snapshot: snapshot, expiresAt: expiresAt, policyRevision: revCopy,
	}
	return nil
}

func (m *memStore) ConsumeActivationToken(ctx context.Context, token string, projectID int64, planHash string) (*WorkflowPlan, error) {
	_, plan, err := m.consumeActivationTokenWithRev(ctx, token, projectID, planHash)
	return plan, err
}

func (m *memStore) consumeActivationTokenLocked(ctx context.Context, token string, projectID int64, planHash string) (*int64, error) {
	rev, _, err := m.consumeActivationTokenWithRev(ctx, token, projectID, planHash)
	return rev, err
}

func (m *memStore) consumeActivationTokenWithRev(_ context.Context, token string, projectID int64, planHash string) (*int64, *WorkflowPlan, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	tokenHash := hashActivationToken(token)
	rec, ok := m.tokens[tokenHash]
	if !ok || rec.used || rec.projectID != projectID || rec.planHash != planHash || time.Now().After(rec.expiresAt) {
		return nil, nil, ErrTokenInvalid
	}
	rec.used = true
	m.tokens[tokenHash] = rec
	plan := rec.snapshot
	return rec.policyRevision, &plan, nil
}