breakpilot-compliance/ai-compliance-sdk/internal/ucca/tom_obligation_mapper.go

package ucca

import (
	"sort"
)

// TOMObligationMapper provides bidirectional mapping between obligations and TOM controls
type TOMObligationMapper struct {
	tomIndex            *TOMControlIndex
	obligationToControl map[string][]string // obligation_id -> []control_id
	controlToObligation map[string][]string // control_id -> []obligation_id
}

// TOMControlRequirement represents a required TOM control with context
type TOMControlRequirement struct {
	Control          *TOMControl `json:"control"`
	ObligationIDs    []string    `json:"obligation_ids"`
	Priority         string      `json:"priority"`         // highest priority from linked obligations
	RequiredByCount  int         `json:"required_by_count"` // number of obligations requiring this
}

// NewTOMObligationMapper creates a new mapper from TOM index and v2 mapping
func NewTOMObligationMapper(tomIndex *TOMControlIndex, mapping *V2TOMMapping) *TOMObligationMapper {
	m := &TOMObligationMapper{
		tomIndex:            tomIndex,
		obligationToControl: make(map[string][]string),
		controlToObligation: make(map[string][]string),
	}

	if mapping != nil {
		m.obligationToControl = mapping.ObligationToControl
		m.controlToObligation = mapping.ControlToObligation
	}

	return m
}

// NewTOMObligationMapperFromObligations builds the mapper from obligations' tom_control_ids
func NewTOMObligationMapperFromObligations(tomIndex *TOMControlIndex, obligations []V2Obligation) *TOMObligationMapper {
	m := &TOMObligationMapper{
		tomIndex:            tomIndex,
		obligationToControl: make(map[string][]string),
		controlToObligation: make(map[string][]string),
	}

	for _, obl := range obligations {
		for _, controlID := range obl.TOMControlIDs {
			m.obligationToControl[obl.ID] = append(m.obligationToControl[obl.ID], controlID)
			m.controlToObligation[controlID] = append(m.controlToObligation[controlID], obl.ID)
		}
	}

	return m
}

// GetControlsForObligation returns TOM controls linked to an obligation
func (m *TOMObligationMapper) GetControlsForObligation(obligationID string) []*TOMControl {
	controlIDs := m.obligationToControl[obligationID]
	var result []*TOMControl
	for _, id := range controlIDs {
		if ctrl, ok := m.tomIndex.GetControl(id); ok {
			result = append(result, ctrl)
		}
	}
	return result
}

// GetControlIDsForObligation returns control IDs for an obligation
func (m *TOMObligationMapper) GetControlIDsForObligation(obligationID string) []string {
	return m.obligationToControl[obligationID]
}

// GetObligationsForControl returns obligation IDs linked to a control
func (m *TOMObligationMapper) GetObligationsForControl(controlID string) []string {
	return m.controlToObligation[controlID]
}

// DeriveControlsFromObligations takes a list of applicable obligations and returns
// deduplicated, priority-sorted TOM control requirements
func (m *TOMObligationMapper) DeriveControlsFromObligations(obligations []Obligation) []TOMControlRequirement {
	// Collect all required controls with their linking obligations
	controlMap := make(map[string]*TOMControlRequirement)

	priorityRank := map[string]int{"critical": 0, "high": 1, "medium": 2, "low": 3}

	for _, obl := range obligations {
		// Get control IDs from ExternalResources (where tom_control_ids are stored)
		controlIDs := obl.ExternalResources
		if len(controlIDs) == 0 {
			controlIDs = m.obligationToControl[obl.ID]
		}

		for _, controlID := range controlIDs {
			ctrl, ok := m.tomIndex.GetControl(controlID)
			if !ok {
				continue
			}

			if existing, found := controlMap[controlID]; found {
				existing.ObligationIDs = append(existing.ObligationIDs, obl.ID)
				existing.RequiredByCount++
				// Keep highest priority
				if rank, ok := priorityRank[string(obl.Priority)]; ok {
					if existingRank, ok2 := priorityRank[existing.Priority]; ok2 && rank < existingRank {
						existing.Priority = string(obl.Priority)
					}
				}
			} else {
				controlMap[controlID] = &TOMControlRequirement{
					Control:         ctrl,
					ObligationIDs:   []string{obl.ID},
					Priority:        string(obl.Priority),
					RequiredByCount: 1,
				}
			}
		}
	}

	// Convert to slice and sort by priority then required_by_count
	var result []TOMControlRequirement
	for _, req := range controlMap {
		result = append(result, *req)
	}

	sort.Slice(result, func(i, j int) bool {
		ri := priorityRank[result[i].Priority]
		rj := priorityRank[result[j].Priority]
		if ri != rj {
			return ri < rj
		}
		return result[i].RequiredByCount > result[j].RequiredByCount
	})

	return result
}

// AddMapping adds a single obligation->control mapping
func (m *TOMObligationMapper) AddMapping(obligationID, controlID string) {
	m.obligationToControl[obligationID] = appendUnique(m.obligationToControl[obligationID], controlID)
	m.controlToObligation[controlID] = appendUnique(m.controlToObligation[controlID], obligationID)
}

func appendUnique(slice []string, item string) []string {
	for _, s := range slice {
		if s == item {
			return slice
		}
	}
	return append(slice, item)
}