breakpilot-compliance/ai-compliance-sdk/internal/api/handlers/iace_handler_enrich.go

package handlers

import (
	"fmt"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
)

// RegulatoryHint represents a relevant passage from TRBS/TRGS/ASR/OSHA.
type RegulatoryHint struct {
	RegulationID    string  `json:"regulation_id"`
	RegulationShort string  `json:"regulation_short"`
	Category        string  `json:"category"`
	Text            string  `json:"text"`
	Pages           []int   `json:"pages,omitempty"`
	SourceURL       string  `json:"source_url,omitempty"`
	Score           float64 `json:"score"`
}

// categoryToSearchTerms maps hazard categories to German search terms
// that match TRBS/TRGS/ASR/OSHA content.
var categoryToSearchTerms = map[string]string{
	"mechanical_hazard":          "mechanische Gefaehrdung Quetschstelle Scherstelle Stossstelle",
	"electrical_hazard":          "elektrische Gefaehrdung Stromschlag Lichtbogen Kurzschluss",
	"thermal_hazard":             "thermische Gefaehrdung Verbrennung Erfrierung heisse Oberflaeche",
	"noise_hazard":               "Laerm Gehoerschutz Schalldruckpegel Laermexposition",
	"vibration_hazard":           "Vibration Hand-Arm Ganzkoerper Schwingungsbelastung",
	"radiation_hazard":           "Strahlung ionisierend nichtionisierend Laser UV",
	"chemical_hazard":            "Gefahrstoff chemische Gefaehrdung Exposition Grenzwert",
	"ergonomic_hazard":           "Ergonomie Zwangshaltung Lasthandhabung Koerperbelastung",
	"hydraulic_hazard":           "Hydraulik Druckbehaelter Druck Bersten Leckage",
	"pneumatic_hazard":           "Pneumatik Druckluft Druckbehaelter Belueftung",
	"software_hazard":            "Software Sicherheitsfunktion Fehlfunktion Programmierung",
	"safety_function_failure":    "Sicherheitsfunktion Ausfall SIL Performance Level",
	"fire_explosion_hazard":      "Brand Explosion explosionsfaehige Atmosphaere Zuendschutz",
	"falling_hazard":             "Absturz herabfallende Gegenstaende Sturzgefahr",
	"trip_slip_hazard":           "Stolpern Rutschen Ausrutschen Fussboden Verkehrsweg",
	"entrapment_hazard":          "Einzugsstelle Fangstelle rotierende Teile Wickelgefahr",
	"crush_hazard":               "Quetschgefahr Quetschstelle Einklemmen Andruckkraft",
	"cut_hazard":                 "Schneiden Schneidwerkzeug Schnittverletzung scharfe Kante",
	"stabbing_hazard":            "Stechen Stichverletzung spitze Teile Injektionsgefahr",
	"high_pressure_hazard":       "Hochdruck Fluessigkeitsstrahl Druckbehaelter Ueberdruck",
	"collision_hazard":           "Kollision Zusammenstoss Anfahren fahrerlose Transportsysteme",
	"lack_of_stability_hazard":   "Standsicherheit Umkippen Kippen Stabilitaet",
	"unexpected_start_hazard":    "unerwarteter Anlauf Wiederanlauf Energietrennung Lockout",
	"control_system_failure":     "Steuerungsausfall Steuerung Fehler Ausfall Sicherheitssteuerung",
	"ppe_hazard":                 "PSA persoenliche Schutzausruestung Schutzkleidung",
}

// EnrichHazardWithRegulations returns regulatory hints for a specific hazard.
// GET /projects/:id/hazards/:hid/regulatory-hints
func (h *IACEHandler) EnrichHazardWithRegulations(c *gin.Context) {
	projectID, err := uuid.Parse(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid project ID"})
		return
	}

	hazardID, err := uuid.Parse(c.Param("hid"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid hazard ID"})
		return
	}

	// Fetch hazard
	hazard, err := h.store.GetHazard(c.Request.Context(), hazardID)
	if err != nil || hazard == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "hazard not found"})
		return
	}
	if hazard.ProjectID != projectID {
		c.JSON(http.StatusNotFound, gin.H{"error": "hazard not in this project"})
		return
	}

	// Fetch project for machine context
	project, err := h.store.GetProject(c.Request.Context(), projectID)
	if err != nil || project == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "project not found"})
		return
	}

	// Build search query from hazard context
	query := buildHazardSearchQuery(hazard.Category, hazard.Name, hazard.Scenario, project.MachineName, project.MachineType)

	// Search bp_compliance_ce (TRBS/TRGS/ASR/OSHA)
	results, err := h.ragClient.SearchCollection(
		c.Request.Context(),
		"bp_compliance_ce",
		query,
		nil, // no regulation filter — search all
		7,
	)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "regulatory search failed: " + err.Error(),
		})
		return
	}

	hints := make([]RegulatoryHint, 0, len(results))
	for _, r := range results {
		if r.Score < 0.3 {
			continue // skip low-relevance results
		}
		hints = append(hints, RegulatoryHint{
			RegulationID:    r.RegulationCode,
			RegulationShort: r.RegulationShort,
			Category:        r.Category,
			Text:            truncateHintText(r.Text, 500),
			Pages:           r.Pages,
			SourceURL:       r.SourceURL,
			Score:           r.Score,
		})
	}

	c.JSON(http.StatusOK, gin.H{
		"hazard_id":   hazardID.String(),
		"hazard_name": hazard.Name,
		"category":    hazard.Category,
		"query":       query,
		"hints":       hints,
		"total":       len(hints),
	})
}

// EnrichMitigationWithRegulations returns regulatory hints for a mitigation.
// GET /projects/:id/mitigations/:mid/regulatory-hints
func (h *IACEHandler) EnrichMitigationWithRegulations(c *gin.Context) {
	projectID, err := uuid.Parse(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid project ID"})
		return
	}

	mitigationID, err := uuid.Parse(c.Param("mid"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid mitigation ID"})
		return
	}

	// Fetch mitigation
	mitigation, err := h.store.GetMitigation(c.Request.Context(), mitigationID)
	if err != nil || mitigation == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "mitigation not found"})
		return
	}

	// Fetch the hazard to get category context
	hazard, err := h.store.GetHazard(c.Request.Context(), mitigation.HazardID)
	if err != nil || hazard == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "linked hazard not found"})
		return
	}
	if hazard.ProjectID != projectID {
		c.JSON(http.StatusNotFound, gin.H{"error": "mitigation not in this project"})
		return
	}

	// Build search query from mitigation + hazard context
	queryParts := []string{mitigation.Name}
	if mitigation.Description != "" {
		queryParts = append(queryParts, mitigation.Description)
	}
	queryParts = append(queryParts, "Schutzmassnahme")
	if terms, ok := categoryToSearchTerms[hazard.Category]; ok {
		queryParts = append(queryParts, terms)
	}
	query := strings.Join(queryParts, " ")
	if len(query) > 500 {
		query = query[:500]
	}

	results, err := h.ragClient.SearchCollection(
		c.Request.Context(),
		"bp_compliance_ce",
		query,
		nil,
		7,
	)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "regulatory search failed: " + err.Error(),
		})
		return
	}

	hints := make([]RegulatoryHint, 0, len(results))
	for _, r := range results {
		if r.Score < 0.3 {
			continue
		}
		hints = append(hints, RegulatoryHint{
			RegulationID:    r.RegulationCode,
			RegulationShort: r.RegulationShort,
			Category:        r.Category,
			Text:            truncateHintText(r.Text, 500),
			Pages:           r.Pages,
			SourceURL:       r.SourceURL,
			Score:           r.Score,
		})
	}

	c.JSON(http.StatusOK, gin.H{
		"mitigation_id":   mitigationID.String(),
		"mitigation_name": mitigation.Name,
		"reduction_type":  mitigation.ReductionType,
		"query":           query,
		"hints":           hints,
		"total":           len(hints),
	})
}

// buildHazardSearchQuery creates a contextual query for RAG search.
func buildHazardSearchQuery(category, name, scenario, machineName, machineType string) string {
	parts := make([]string, 0, 5)

	// Add category-specific German search terms
	if terms, ok := categoryToSearchTerms[category]; ok {
		parts = append(parts, terms)
	}

	// Add hazard name and scenario
	if name != "" {
		parts = append(parts, name)
	}
	if scenario != "" && len(scenario) < 200 {
		parts = append(parts, scenario)
	}

	// Add machine context
	if machineType != "" {
		parts = append(parts, machineType)
	}
	if machineName != "" && len(parts) < 4 {
		parts = append(parts, machineName)
	}

	query := strings.Join(parts, " ")
	if len(query) > 500 {
		query = query[:500]
	}
	return query
}

func truncateHintText(text string, maxLen int) string {
	if len(text) <= maxLen {
		return text
	}
	// Find last sentence boundary
	truncated := text[:maxLen]
	if lastDot := strings.LastIndex(truncated, ". "); lastDot > maxLen/2 {
		return truncated[:lastDot+1]
	}
	return truncated + "..."
}

// ============================================================================
// Batch: Enrich all hazards at once (for overview display)
// ============================================================================

// EnrichProjectHazardsBatch returns top regulatory hint per hazard category.
// GET /projects/:id/regulatory-hints
func (h *IACEHandler) EnrichProjectHazardsBatch(c *gin.Context) {
	projectID, err := uuid.Parse(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid project ID"})
		return
	}

	project, err := h.store.GetProject(c.Request.Context(), projectID)
	if err != nil || project == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "project not found"})
		return
	}

	// Get all hazards to extract unique categories
	hazards, err := h.store.ListHazards(c.Request.Context(), projectID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list hazards"})
		return
	}

	// Deduplicate categories
	seen := make(map[string]bool)
	var categories []string
	for _, hz := range hazards {
		if !seen[hz.Category] {
			seen[hz.Category] = true
			categories = append(categories, hz.Category)
		}
	}

	// One RAG search per unique category (typically 5-10 categories, not 160 hazards)
	type CategoryHints struct {
		Category string           `json:"category"`
		Hints    []RegulatoryHint `json:"hints"`
	}

	result := make([]CategoryHints, 0, len(categories))
	for _, cat := range categories {
		query := buildHazardSearchQuery(cat, "", "", project.MachineName, project.MachineType)
		results, err := h.ragClient.SearchCollection(
			c.Request.Context(),
			"bp_compliance_ce",
			query,
			nil,
			3,
		)
		if err != nil {
			continue
		}

		hints := make([]RegulatoryHint, 0, len(results))
		for _, r := range results {
			if r.Score < 0.3 {
				continue
			}
			hints = append(hints, RegulatoryHint{
				RegulationID:    r.RegulationCode,
				RegulationShort: r.RegulationShort,
				Category:        r.Category,
				Text:            truncateHintText(r.Text, 300),
				Pages:           r.Pages,
				SourceURL:       r.SourceURL,
				Score:           r.Score,
			})
		}
		if len(hints) > 0 {
			result = append(result, CategoryHints{Category: cat, Hints: hints})
		}
	}

	c.JSON(http.StatusOK, gin.H{
		"project_id":       projectID.String(),
		"categories":       len(categories),
		"total_hazards":    len(hazards),
		"regulatory_hints": result,
		"sources":          fmt.Sprintf("TRBS/TRGS/ASR (%d BAuA) + OSHA Technical Manual", 126),
	})
}