Agent-System/internal/traceimpact/store.go

package traceimpact

import (
	"context"
	"fmt"
	"sync"
)

// Store persists trace rows. Implementations must not delete decision history.
type Store interface {
	UpsertItem(ctx context.Context, item Item) (Item, error)
	ListItems(ctx context.Context, projectID int64) ([]Item, error)
	UpsertEdge(ctx context.Context, edge Edge) (Edge, error)
	ListEdges(ctx context.Context, projectID int64) ([]Edge, error)
	PutIndex(ctx context.Context, st IndexState) error
	IndexStatus(ctx context.Context, projectID int64, docID string) (IndexState, bool, error)
	ListIndex(ctx context.Context, projectID int64) ([]IndexState, error)
	SetEdgeReview(ctx context.Context, edgeID int64, review string) error
	AddDecision(ctx context.Context, d Decision) (Decision, error)
	ListDecisions(ctx context.Context, projectID int64) ([]Decision, error)
	PutReopen(ctx context.Context, req ReopenRequest) (ReopenRequest, error)
	GetReopen(ctx context.Context, id int64) (ReopenRequest, error)
	ListPendingReopens(ctx context.Context, projectID int64) ([]ReopenRequest, error)
	PutMark(ctx context.Context, m RecheckMark) (RecheckMark, error)
	ListOpenMarks(ctx context.Context, projectID int64, stageKey string) ([]RecheckMark, error)
	AckMark(ctx context.Context, id, userID int64) error
	SetMarkStatus(ctx context.Context, id int64, status, reason string) error
	GetEdge(ctx context.Context, edgeID int64) (Edge, error)
	AppendRevision(ctx context.Context, rev ItemRevision) error
	RevisionCount(ctx context.Context, itemID int64) (int, error)
	AddProposal(ctx context.Context, projectID, userID int64, text string) (int64, error)
}

type memStore struct {
	mu        sync.Mutex
	items     map[int64]Item
	byKey     map[string]int64
	edges     map[int64]Edge
	byEdge    map[string]int64
	index     map[string]IndexState
	decisions []Decision
	reopens   map[int64]ReopenRequest
	marks     map[int64]RecheckMark
	revs      map[int64][]ItemRevision
	props     []proposalRow
	nextItem  int64
	nextEdge  int64
	nextDec   int64
	nextReq   int64
	nextMark  int64
	nextProp  int64
}

func newMemStore() *memStore {
	return &memStore{
		items:   map[int64]Item{},
		byKey:   map[string]int64{},
		edges:   map[int64]Edge{},
		byEdge:  map[string]int64{},
		index:   map[string]IndexState{},
		reopens: map[int64]ReopenRequest{},
		marks:   map[int64]RecheckMark{},
	}
}

func itemKey(projectID int64, kind, business string) string {
	return fmt.Sprintf("%d|%s|%s", projectID, kind, business)
}

type proposalRow struct {
	id        int64
	projectID int64
	userID    int64
	text      string
}

func (m *memStore) UpsertItem(_ context.Context, item Item) (Item, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	k := itemKey(item.ProjectID, item.Kind, item.BusinessKey)
	if id, ok := m.byKey[k]; ok {
		cur := m.items[id]
		if item.Statement != "" {
			cur.Statement = item.Statement
		}
		if item.SourceDocID != "" {
			cur.SourceDocID = item.SourceDocID
		}
		cur.ReviewState = item.ReviewState
		cur.Tombstone = item.Tombstone
		m.items[id] = cur
		return cur, nil
	}
	m.nextItem++
	item.ID = m.nextItem
	if item.Lifecycle == "" {
		item.Lifecycle = "active"
	}
	m.items[item.ID] = item
	m.byKey[k] = item.ID
	return item, nil
}

func (m *memStore) ListItems(_ context.Context, projectID int64) ([]Item, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []Item
	for _, it := range m.items {
		if it.ProjectID == projectID {
			out = append(out, it)
		}
	}
	return out, nil
}

func (m *memStore) UpsertEdge(_ context.Context, edge Edge) (Edge, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	k := fmt.Sprintf("%d|%d|%d|%s|%s", edge.ProjectID, edge.FromItemID, edge.ToItemID, edge.RelationType, edge.Provenance)
	if id, ok := m.byEdge[k]; ok {
		cur := m.edges[id]
		if cur.ReviewState == ReviewRejected && edge.ReviewState == ReviewCandidate {
			return cur, nil
		}
		if edge.ReviewState != "" && cur.ReviewState != ReviewConfirmed {
			cur.ReviewState = edge.ReviewState
		}
		m.edges[id] = cur
		return cur, nil
	}
	m.nextEdge++
	edge.ID = m.nextEdge
	m.edges[edge.ID] = edge
	m.byEdge[k] = edge.ID
	return edge, nil
}

func (m *memStore) ListEdges(_ context.Context, projectID int64) ([]Edge, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []Edge
	for _, e := range m.edges {
		if e.ProjectID == projectID {
			out = append(out, e)
		}
	}
	return out, nil
}

func (m *memStore) PutIndex(_ context.Context, st IndexState) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if st.Extractor == "" {
		st.Extractor = ExtractorVersion
	}
	m.index[fmt.Sprintf("%d|%s", st.ProjectID, st.DocID)] = st
	return nil
}

func (m *memStore) IndexStatus(_ context.Context, projectID int64, docID string) (IndexState, bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	st, ok := m.index[fmt.Sprintf("%d|%s", projectID, docID)]
	return st, ok, nil
}

func (m *memStore) ListIndex(_ context.Context, projectID int64) ([]IndexState, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []IndexState
	for _, st := range m.index {
		if st.ProjectID == projectID {
			out = append(out, st)
		}
	}
	return out, nil
}

func (m *memStore) SetEdgeReview(_ context.Context, edgeID int64, review string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	e, ok := m.edges[edgeID]
	if !ok {
		return fmt.Errorf("traceimpact: edge not found")
	}
	e.ReviewState = review
	m.edges[edgeID] = e
	return nil
}

func (m *memStore) AddDecision(_ context.Context, d Decision) (Decision, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.nextDec++
	d.ID = m.nextDec
	m.decisions = append(m.decisions, d)
	return d, nil
}

func (m *memStore) ListDecisions(_ context.Context, projectID int64) ([]Decision, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []Decision
	for _, d := range m.decisions {
		if d.ProjectID == projectID {
			out = append(out, d)
		}
	}
	return out, nil
}

func (m *memStore) PutReopen(_ context.Context, req ReopenRequest) (ReopenRequest, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if req.ID == 0 {
		m.nextReq++
		req.ID = m.nextReq
	}
	m.reopens[req.ID] = req
	return req, nil
}

func (m *memStore) GetReopen(_ context.Context, id int64) (ReopenRequest, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	req, ok := m.reopens[id]
	if !ok {
		return ReopenRequest{}, fmt.Errorf("traceimpact: reopen request not found")
	}
	return req, nil
}

func (m *memStore) ListPendingReopens(_ context.Context, projectID int64) ([]ReopenRequest, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []ReopenRequest
	for _, req := range m.reopens {
		if req.ProjectID == projectID && req.Status == ReopenPending {
			out = append(out, req)
		}
	}
	return out, nil
}

func (m *memStore) PutMark(_ context.Context, mark RecheckMark) (RecheckMark, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, existing := range m.marks {
		if existing.ProjectID == mark.ProjectID && existing.StageKey == mark.StageKey && existing.SourceDoc == mark.SourceDoc && existing.Status == MarkOpen {
			return existing, nil
		}
	}
	m.nextMark++
	mark.ID = m.nextMark
	if mark.Status == "" {
		mark.Status = MarkOpen
	}
	m.marks[mark.ID] = mark
	return mark, nil
}

func (m *memStore) ListOpenMarks(_ context.Context, projectID int64, stageKey string) ([]RecheckMark, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []RecheckMark
	for _, mark := range m.marks {
		if mark.ProjectID != projectID || mark.Status != MarkOpen {
			continue
		}
		if stageKey != "" && mark.StageKey != stageKey {
			continue
		}
		out = append(out, mark)
	}
	return out, nil
}

func (m *memStore) AckMark(_ context.Context, id, userID int64) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	mark, ok := m.marks[id]
	if !ok {
		return fmt.Errorf("traceimpact: mark not found")
	}
	mark.AckedBy = userID
	m.marks[id] = mark
	return nil
}

func (m *memStore) SetMarkStatus(_ context.Context, id int64, status, reason string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	mark, ok := m.marks[id]
	if !ok {
		return fmt.Errorf("traceimpact: mark not found")
	}
	mark.Status = status
	if reason != "" {
		mark.Reason = reason
	}
	m.marks[id] = mark
	return nil
}

func (m *memStore) GetEdge(_ context.Context, edgeID int64) (Edge, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	e, ok := m.edges[edgeID]
	if !ok {
		return Edge{}, fmt.Errorf("traceimpact: edge not found")
	}
	return e, nil
}

func (m *memStore) AppendRevision(_ context.Context, rev ItemRevision) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.revs == nil {
		m.revs = map[int64][]ItemRevision{}
	}
	for _, existing := range m.revs[rev.ItemID] {
		if existing.SnapshotID == rev.SnapshotID && existing.ContentHash == rev.ContentHash && existing.Extractor == rev.Extractor {
			return nil
		}
	}
	rev.RevisionNo = len(m.revs[rev.ItemID]) + 1
	m.revs[rev.ItemID] = append(m.revs[rev.ItemID], rev)
	return nil
}

func (m *memStore) RevisionCount(_ context.Context, itemID int64) (int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.revs[itemID]), nil
}

func (m *memStore) AddProposal(_ context.Context, projectID, userID int64, text string) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.nextProp++
	m.props = append(m.props, proposalRow{id: m.nextProp, projectID: projectID, userID: userID, text: text})
	return m.nextProp, nil
}