breakpilot-compliance/ai-compliance-sdk/internal/iace/hazard_blocks.go

package iace

import (
	"sort"
	"strings"

	"github.com/google/uuid"
)

// HazardBlock groups related hazards under a parent hazard.
// The parent is the hazard with the highest inherent risk in the group.
// Child hazards are covered by the same or similar protective measures.
type HazardBlock struct {
	ParentHazard HazardBlockEntry   `json:"parent_hazard"`
	Children     []HazardBlockEntry `json:"children"`
	BlockKey     string             `json:"block_key"`
	SharedMeasureCount int          `json:"shared_measure_count"`
	// If true, the parent's measures cover all children → children
	// don't need individual risk assessment.
	ChildrenCoveredByParent bool `json:"children_covered_by_parent"`
}

// HazardBlockEntry is a hazard with its assessment and linked measures.
type HazardBlockEntry struct {
	Hazard         Hazard          `json:"hazard"`
	Assessment     *RiskAssessment `json:"assessment,omitempty"`
	MitigationIDs  []uuid.UUID     `json:"mitigation_ids"`
}

// ComputeHazardBlocks groups hazards into blocks based on category + component.
// Within each block, the hazard with the highest risk becomes the parent.
// Children whose measures are a subset of the parent's measures are marked as covered.
func ComputeHazardBlocks(
	hazards []Hazard,
	assessments []RiskAssessment,
	mitigations []Mitigation,
) []HazardBlock {
	if len(hazards) == 0 {
		return nil
	}

	// Build assessment lookup: hazard_id → latest assessment
	assessMap := make(map[uuid.UUID]*RiskAssessment)
	for i := range assessments {
		a := &assessments[i]
		if existing, ok := assessMap[a.HazardID]; !ok || a.Version > existing.Version {
			assessMap[a.HazardID] = a
		}
	}

	// Build mitigation lookup: hazard_id → []mitigation_ids
	mitsByHazard := make(map[uuid.UUID][]uuid.UUID)
	for _, m := range mitigations {
		mitsByHazard[m.HazardID] = append(mitsByHazard[m.HazardID], m.ID)
	}

	// Group by blockKey = category + ":" + componentID
	groups := make(map[string][]HazardBlockEntry)
	for _, h := range hazards {
		key := buildBlockKey(h)
		entry := HazardBlockEntry{
			Hazard:        h,
			Assessment:    assessMap[h.ID],
			MitigationIDs: mitsByHazard[h.ID],
		}
		groups[key] = append(groups[key], entry)
	}

	// Build blocks: sort each group by risk, first is parent
	var blocks []HazardBlock
	for key, entries := range groups {
		sortByRiskDesc(entries, assessMap)

		parent := entries[0]
		children := entries[1:]

		// Check if parent's measures cover children
		parentMitSet := toUUIDSet(parent.MitigationIDs)
		allCovered := true
		for _, child := range children {
			if !mitigationsCoveredBy(child, parent, mitigations) {
				allCovered = false
				break
			}
		}

		block := HazardBlock{
			ParentHazard:            parent,
			Children:                children,
			BlockKey:                key,
			SharedMeasureCount:      len(parentMitSet),
			ChildrenCoveredByParent: allCovered && len(children) > 0,
		}
		blocks = append(blocks, block)
	}

	// Sort blocks: largest (most children) first, then by parent risk
	sort.Slice(blocks, func(i, j int) bool {
		ri := inherentRisk(blocks[i].ParentHazard, assessMap)
		rj := inherentRisk(blocks[j].ParentHazard, assessMap)
		if len(blocks[i].Children) != len(blocks[j].Children) {
			return len(blocks[i].Children) > len(blocks[j].Children)
		}
		return ri > rj
	})

	return blocks
}

func buildBlockKey(h Hazard) string {
	comp := h.ComponentID.String()
	zone := NormalizeDEPublic(h.HazardousZone)
	// Use first 3 significant words of zone for grouping
	words := strings.Fields(zone)
	var sig []string
	for _, w := range words {
		w = strings.Trim(w, ".,;:!?()/-")
		if len(w) >= 4 {
			sig = append(sig, w)
		}
		if len(sig) >= 2 {
			break
		}
	}
	zoneKey := strings.Join(sig, "_")
	return h.Category + ":" + comp + ":" + zoneKey
}

func sortByRiskDesc(entries []HazardBlockEntry, assessMap map[uuid.UUID]*RiskAssessment) {
	sort.Slice(entries, func(i, j int) bool {
		ri := inherentRisk(entries[i], assessMap)
		rj := inherentRisk(entries[j], assessMap)
		return ri > rj
	})
}

func inherentRisk(entry HazardBlockEntry, assessMap map[uuid.UUID]*RiskAssessment) float64 {
	if entry.Assessment != nil {
		return entry.Assessment.InherentRisk
	}
	if a, ok := assessMap[entry.Hazard.ID]; ok {
		return a.InherentRisk
	}
	return 0
}

// mitigationsCoveredBy checks if child's measures are functionally covered
// by parent's measures (same reduction type and hazard category).
func mitigationsCoveredBy(child, parent HazardBlockEntry, allMits []Mitigation) bool {
	if len(child.MitigationIDs) == 0 {
		return true // No measures needed → covered by default
	}

	mitMap := make(map[uuid.UUID]Mitigation)
	for _, m := range allMits {
		mitMap[m.ID] = m
	}

	// Check: for each child mitigation type, parent has same type
	parentTypes := make(map[ReductionType]bool)
	for _, mid := range parent.MitigationIDs {
		if m, ok := mitMap[mid]; ok {
			parentTypes[m.ReductionType] = true
		}
	}

	for _, mid := range child.MitigationIDs {
		if m, ok := mitMap[mid]; ok {
			if !parentTypes[m.ReductionType] {
				return false
			}
		}
	}
	return true
}

func toUUIDSet(ids []uuid.UUID) map[uuid.UUID]bool {
	s := make(map[uuid.UUID]bool, len(ids))
	for _, id := range ids {
		s[id] = true
	}
	return s
}