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

package iace

import (
	"strings"
	"testing"

	"github.com/google/uuid"
)

// ============================================================================
// Test Helpers
// ============================================================================

// newTestSectionContext builds a SectionGenerationContext with realistic sample
// data for a "Robot Arm XY-200" industrial robot project.
func newTestSectionContext() *SectionGenerationContext {
	projectID := uuid.New()
	compSoftwareID := uuid.New()
	compSensorID := uuid.New()
	hazardHighID := uuid.New()
	hazardLowID := uuid.New()

	return &SectionGenerationContext{
		Project: &Project{
			ID:              projectID,
			TenantID:        uuid.New(),
			MachineName:     "Robot Arm XY-200",
			MachineType:     "industrial_robot",
			Manufacturer:    "TestCorp",
			Description:     "6-axis industrial robot for automotive welding",
			CEMarkingTarget: "2023/1230",
			Status:          ProjectStatusHazardAnalysis,
		},
		Components: []Component{
			{
				ID:               compSoftwareID,
				ProjectID:        projectID,
				Name:             "SafetyPLC-500",
				ComponentType:    ComponentTypeSoftware,
				Version:          "3.2.1",
				Description:      "Safety-rated programmable logic controller firmware",
				IsSafetyRelevant: true,
				IsNetworked:      true,
			},
			{
				ID:               compSensorID,
				ProjectID:        projectID,
				Name:             "ProxSensor-LiDAR",
				ComponentType:    ComponentTypeSensor,
				Version:          "1.0.0",
				Description:      "LiDAR proximity sensor for collision avoidance",
				IsSafetyRelevant: false,
				IsNetworked:      false,
			},
		},
		Hazards: []Hazard{
			{
				ID:             hazardHighID,
				ProjectID:      projectID,
				ComponentID:    compSoftwareID,
				Name:           "Software malfunction causing uncontrolled movement",
				Description:    "Firmware fault leads to unpredictable arm trajectory",
				Category:       "mechanical",
				SubCategory:    "crushing",
				Status:         HazardStatusAssessed,
				LifecyclePhase: "normal_operation",
				AffectedPerson: "operator",
				PossibleHarm:   "Severe crushing injury to operator",
			},
			{
				ID:             hazardLowID,
				ProjectID:      projectID,
				ComponentID:    compSensorID,
				Name:           "Sensor drift causing delayed stop",
				Description:    "Gradual sensor calibration loss reduces reaction time",
				Category:       "electrical",
				SubCategory:    "sensor_failure",
				Status:         HazardStatusIdentified,
				LifecyclePhase: "normal_operation",
				AffectedPerson: "bystander",
				PossibleHarm:   "Minor bruising",
			},
		},
		Assessments: map[uuid.UUID][]RiskAssessment{
			hazardHighID: {
				{
					ID:             uuid.New(),
					HazardID:       hazardHighID,
					Version:        1,
					AssessmentType: AssessmentTypeInitial,
					Severity:       5,
					Exposure:       4,
					Probability:    3,
					Avoidance:      2,
					InherentRisk:   120.0,
					ResidualRisk:   85.0,
					RiskLevel:      RiskLevelHigh,
					IsAcceptable:   false,
				},
			},
			hazardLowID: {
				{
					ID:             uuid.New(),
					HazardID:       hazardLowID,
					Version:        1,
					AssessmentType: AssessmentTypeInitial,
					Severity:       2,
					Exposure:       3,
					Probability:    2,
					Avoidance:      4,
					InherentRisk:   12.0,
					ResidualRisk:   6.0,
					RiskLevel:      RiskLevelLow,
					IsAcceptable:   true,
				},
			},
		},
		Mitigations: map[uuid.UUID][]Mitigation{
			hazardHighID: {
				{
					ID:            uuid.New(),
					HazardID:      hazardHighID,
					ReductionType: ReductionTypeDesign,
					Name:          "Redundant safety controller",
					Description:   "Dual-channel safety PLC with cross-monitoring",
					Status:        MitigationStatusImplemented,
				},
				{
					ID:            uuid.New(),
					HazardID:      hazardHighID,
					ReductionType: ReductionTypeProtective,
					Name:          "Light curtain barrier",
					Description:   "Type 4 safety light curtain around work envelope",
					Status:        MitigationStatusVerified,
				},
			},
			hazardLowID: {
				{
					ID:            uuid.New(),
					HazardID:      hazardLowID,
					ReductionType: ReductionTypeInformation,
					Name:          "Calibration schedule warning",
					Description:   "Automated alert when sensor calibration is overdue",
					Status:        MitigationStatusPlanned,
				},
			},
		},
		Classifications: []RegulatoryClassification{
			{
				ID:                   uuid.New(),
				ProjectID:            projectID,
				Regulation:           RegulationMachineryRegulation,
				ClassificationResult: "Annex I - High-Risk Machinery",
				RiskLevel:            RiskLevelHigh,
				Confidence:           0.92,
			},
		},
		Evidence: []Evidence{
			{
				ID:          uuid.New(),
				ProjectID:   projectID,
				FileName:    "safety_plc_test_report.pdf",
				Description: "Functional safety test report for SafetyPLC-500",
			},
		},
		RAGContext: "",
	}
}

// ============================================================================
// Tests: buildUserPrompt
// ============================================================================

func TestBuildUserPrompt_RiskAssessmentReport(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "risk_assessment_report")

	// Must contain project identification
	if !strings.Contains(prompt, "Robot Arm XY-200") {
		t.Error("prompt should contain machine name 'Robot Arm XY-200'")
	}
	if !strings.Contains(prompt, "TestCorp") {
		t.Error("prompt should contain manufacturer 'TestCorp'")
	}

	// Must reference hazard information
	if !strings.Contains(prompt, "uncontrolled movement") || !strings.Contains(prompt, "Software malfunction") {
		t.Error("prompt should contain hazard name or description for high-risk hazard")
	}

	// Must reference risk levels
	if !strings.Contains(prompt, "high") && !strings.Contains(prompt, "High") && !strings.Contains(prompt, "HIGH") {
		t.Error("prompt should reference the high risk level")
	}
}

func TestBuildUserPrompt_ComponentList(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "component_list")

	// Must list component names
	if !strings.Contains(prompt, "SafetyPLC-500") {
		t.Error("prompt should contain component name 'SafetyPLC-500'")
	}
	if !strings.Contains(prompt, "ProxSensor-LiDAR") {
		t.Error("prompt should contain component name 'ProxSensor-LiDAR'")
	}

	// Must reference component types
	if !strings.Contains(prompt, "software") && !strings.Contains(prompt, "Software") {
		t.Error("prompt should contain component type 'software'")
	}
	if !strings.Contains(prompt, "sensor") && !strings.Contains(prompt, "Sensor") {
		t.Error("prompt should contain component type 'sensor'")
	}
}

func TestBuildUserPrompt_EmptyProject(t *testing.T) {
	sctx := &SectionGenerationContext{
		Project:         nil,
		Components:      nil,
		Hazards:         nil,
		Assessments:     nil,
		Mitigations:     nil,
		Classifications: nil,
		Evidence:        nil,
		RAGContext:       "",
	}

	// Should not panic on nil/empty data
	prompt := buildUserPrompt(sctx, "general_description")
	if prompt == "" {
		t.Error("buildUserPrompt should return non-empty string even for empty context")
	}
}

func TestBuildUserPrompt_MitigationReport(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "mitigation_report")

	// Must reference mitigation names
	if !strings.Contains(prompt, "Redundant safety controller") {
		t.Error("prompt should contain design mitigation 'Redundant safety controller'")
	}
	if !strings.Contains(prompt, "Light curtain barrier") {
		t.Error("prompt should contain protective mitigation 'Light curtain barrier'")
	}
	if !strings.Contains(prompt, "Calibration schedule warning") {
		t.Error("prompt should contain information mitigation 'Calibration schedule warning'")
	}

	// Must reference reduction types
	hasDesign := strings.Contains(prompt, "design") || strings.Contains(prompt, "Design")
	hasProtective := strings.Contains(prompt, "protective") || strings.Contains(prompt, "Protective")
	hasInformation := strings.Contains(prompt, "information") || strings.Contains(prompt, "Information")
	if !hasDesign {
		t.Error("prompt should reference 'design' reduction type")
	}
	if !hasProtective {
		t.Error("prompt should reference 'protective' reduction type")
	}
	if !hasInformation {
		t.Error("prompt should reference 'information' reduction type")
	}
}

func TestBuildUserPrompt_WithRAGContext(t *testing.T) {
	sctx := newTestSectionContext()
	sctx.RAGContext = "According to EN ISO 13849-1:2023, safety-related parts of control systems for machinery shall be designed and constructed using the principles of EN ISO 12100."

	prompt := buildUserPrompt(sctx, "standards_applied")

	if !strings.Contains(prompt, "EN ISO 13849-1") {
		t.Error("prompt should include the RAG context referencing EN ISO 13849-1")
	}
	if !strings.Contains(prompt, "EN ISO 12100") {
		t.Error("prompt should include the RAG context referencing EN ISO 12100")
	}
}

func TestBuildUserPrompt_WithoutRAGContext(t *testing.T) {
	sctx := newTestSectionContext()
	sctx.RAGContext = ""

	prompt := buildUserPrompt(sctx, "standards_applied")

	// Should still produce a valid prompt without RAG context
	if prompt == "" {
		t.Error("prompt should be non-empty even without RAG context")
	}
	// Should still contain the project info
	if !strings.Contains(prompt, "Robot Arm XY-200") {
		t.Error("prompt should still contain machine name when no RAG context")
	}
}

// ============================================================================
// Tests: buildRAGQuery
// ============================================================================

func TestBuildRAGQuery_AllSectionTypes(t *testing.T) {
	sectionTypes := []string{
		"risk_assessment_report",
		"hazard_log_combined",
		"general_description",
		"essential_requirements",
		"design_specifications",
		"test_reports",
		"standards_applied",
		"declaration_of_conformity",
		"ai_intended_purpose",
		"ai_model_description",
		"ai_risk_management",
		"ai_human_oversight",
		"component_list",
		"classification_report",
		"mitigation_report",
		"verification_report",
		"evidence_index",
		"instructions_for_use",
		"monitoring_plan",
	}

	for _, st := range sectionTypes {
		t.Run(st, func(t *testing.T) {
			q := buildRAGQuery(st)
			if q == "" {
				t.Errorf("buildRAGQuery(%q) returned empty string", st)
			}
			// Each query should be at least a few words long to be useful
			if len(q) < 10 {
				t.Errorf("buildRAGQuery(%q) returned suspiciously short query: %q", st, q)
			}
		})
	}
}

func TestBuildRAGQuery_UnknownSectionType(t *testing.T) {
	q := buildRAGQuery("nonexistent_section_type")
	// Should return a generic fallback query rather than an empty string
	// (the function needs some query to send to RAG even for unknown types)
	if q == "" {
		t.Log("buildRAGQuery returned empty for unknown section type (may be acceptable if caller handles this)")
	}
}

func TestBuildRAGQuery_QueriesAreDifferent(t *testing.T) {
	// Different section types should produce different queries for targeted retrieval
	q1 := buildRAGQuery("risk_assessment_report")
	q2 := buildRAGQuery("declaration_of_conformity")
	q3 := buildRAGQuery("monitoring_plan")

	if q1 == q2 {
		t.Error("risk_assessment_report and declaration_of_conformity should have different RAG queries")
	}
	if q2 == q3 {
		t.Error("declaration_of_conformity and monitoring_plan should have different RAG queries")
	}
	if q1 == q3 {
		t.Error("risk_assessment_report and monitoring_plan should have different RAG queries")
	}
}

// ============================================================================
// Tests: sectionSystemPrompts
// ============================================================================

func TestSectionSystemPrompts_Coverage(t *testing.T) {
	requiredTypes := []string{
		"risk_assessment_report",
		"hazard_log_combined",
		"general_description",
		"essential_requirements",
		"design_specifications",
		"test_reports",
		"standards_applied",
		"declaration_of_conformity",
		"component_list",
		"classification_report",
		"mitigation_report",
		"verification_report",
		"evidence_index",
		"instructions_for_use",
		"monitoring_plan",
	}

	for _, st := range requiredTypes {
		t.Run(st, func(t *testing.T) {
			prompt, ok := sectionSystemPrompts[st]
			if !ok {
				t.Errorf("sectionSystemPrompts missing entry for %q", st)
				return
			}
			if prompt == "" {
				t.Errorf("sectionSystemPrompts[%q] is empty", st)
			}
			// System prompts should contain meaningful instruction text
			if len(prompt) < 50 {
				t.Errorf("sectionSystemPrompts[%q] is suspiciously short (%d chars)", st, len(prompt))
			}
		})
	}
}

func TestSectionSystemPrompts_ContainRoleInstruction(t *testing.T) {
	// Each system prompt should instruct the LLM about its role as a compliance expert.
	// Prompts are in German, so check for both German and English keywords.
	keywords := []string{
		"expert", "engineer", "compliance", "technical", "documentation", "safety",
		"generate", "write", "create", "produce",
		// German equivalents
		"experte", "ingenieur", "erstelle", "beschreibe", "dokumentation", "sicherheit",
		"risikobeurteilung", "konformit", "norm", "verordnung", "richtlinie",
		"gefaehrdung", "massnahm", "verifikation", "uebersicht", "protokoll",
		"abschnitt", "bericht", "maschin",
	}
	for st, prompt := range sectionSystemPrompts {
		t.Run(st, func(t *testing.T) {
			lower := strings.ToLower(prompt)
			found := false
			for _, kw := range keywords {
				if strings.Contains(lower, kw) {
					found = true
					break
				}
			}
			if !found {
				t.Errorf("sectionSystemPrompts[%q] does not appear to contain a role or task instruction", st)
			}
		})
	}
}

// ============================================================================
// Tests: buildUserPrompt — additional section types
// ============================================================================

func TestBuildUserPrompt_ClassificationReport(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "classification_report")

	// Must reference classification data
	if !strings.Contains(prompt, "machinery_regulation") && !strings.Contains(prompt, "Machinery") && !strings.Contains(prompt, "machinery") {
		t.Error("prompt should reference the machinery regulation classification")
	}
}

func TestBuildUserPrompt_EvidenceIndex(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "evidence_index")

	// Must reference evidence files
	if !strings.Contains(prompt, "safety_plc_test_report.pdf") {
		t.Error("prompt should reference evidence file name 'safety_plc_test_report.pdf'")
	}
}

func TestBuildUserPrompt_GeneralDescription(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "general_description")

	// Must contain machine description
	if !strings.Contains(prompt, "Robot Arm XY-200") {
		t.Error("prompt should contain machine name")
	}
	if !strings.Contains(prompt, "industrial_robot") && !strings.Contains(prompt, "industrial robot") {
		t.Error("prompt should contain machine type")
	}
	if !strings.Contains(prompt, "automotive welding") && !strings.Contains(prompt, "6-axis") {
		t.Error("prompt should reference machine description content")
	}
}

func TestBuildUserPrompt_DeclarationOfConformity(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "declaration_of_conformity")

	// Declaration needs manufacturer and CE target
	if !strings.Contains(prompt, "TestCorp") {
		t.Error("prompt should contain manufacturer for declaration of conformity")
	}
	if !strings.Contains(prompt, "2023/1230") && !strings.Contains(prompt, "CE") && !strings.Contains(prompt, "ce_marking") {
		t.Error("prompt should reference CE marking target or regulation for declaration")
	}
}

// ============================================================================
// Tests: Error handling & edge cases
// ============================================================================

func TestBuildUserPrompt_NilContext(t *testing.T) {
	// Should not panic on completely nil context
	prompt := buildUserPrompt(nil, "risk_assessment_report")
	if prompt == "" {
		t.Error("buildUserPrompt should return non-empty string for nil context")
	}
	if !strings.Contains(prompt, "Keine Projektdaten") {
		t.Error("nil context prompt should contain fallback text 'Keine Projektdaten'")
	}
}

func TestBuildUserPrompt_NilProject(t *testing.T) {
	sctx := &SectionGenerationContext{
		Project: nil,
	}
	prompt := buildUserPrompt(sctx, "general_description")
	if prompt == "" {
		t.Error("buildUserPrompt should return non-empty string for nil project")
	}
	if !strings.Contains(prompt, "Keine Projektdaten") {
		t.Error("nil project prompt should contain fallback text 'Keine Projektdaten'")
	}
}

func TestBuildUserPrompt_EmptyComponents(t *testing.T) {
	sctx := &SectionGenerationContext{
		Project:    &Project{MachineName: "Test", Manufacturer: "Corp"},
		Components: []Component{},
		Hazards:    nil,
	}
	prompt := buildUserPrompt(sctx, "component_list")
	if !strings.Contains(prompt, "Test") {
		t.Error("prompt should contain machine name even with empty components")
	}
}

func TestBuildUserPrompt_EmptyHazards(t *testing.T) {
	sctx := &SectionGenerationContext{
		Project: &Project{MachineName: "Test", Manufacturer: "Corp"},
		Hazards: []Hazard{},
	}
	prompt := buildUserPrompt(sctx, "hazard_log_combined")
	if prompt == "" {
		t.Error("prompt should be non-empty even with no hazards")
	}
}

func TestBuildUserPrompt_EmptyClassifications(t *testing.T) {
	sctx := &SectionGenerationContext{
		Project:         &Project{MachineName: "Test", Manufacturer: "Corp"},
		Classifications: []RegulatoryClassification{},
	}
	prompt := buildUserPrompt(sctx, "classification_report")
	if prompt == "" {
		t.Error("prompt should be non-empty even with no classifications")
	}
}

func TestBuildUserPrompt_EmptyEvidence(t *testing.T) {
	sctx := &SectionGenerationContext{
		Project:  &Project{MachineName: "Test", Manufacturer: "Corp"},
		Evidence: []Evidence{},
	}
	prompt := buildUserPrompt(sctx, "evidence_index")
	if prompt == "" {
		t.Error("prompt should be non-empty even with no evidence")
	}
}

func TestGetSystemPrompt_UnknownType(t *testing.T) {
	prompt := getSystemPrompt("totally_unknown_section")
	if prompt == "" {
		t.Error("getSystemPrompt should return a fallback for unknown types")
	}
	// Fallback should still be a useful CE expert prompt
	lower := strings.ToLower(prompt)
	if !strings.Contains(lower, "ce") && !strings.Contains(lower, "experte") && !strings.Contains(lower, "dokumentation") {
		t.Error("fallback system prompt should reference CE or documentation expertise")
	}
}

func TestGetSystemPrompt_AllKnownTypes(t *testing.T) {
	knownTypes := []string{
		"risk_assessment_report", "hazard_log_combined", "general_description",
		"essential_requirements", "design_specifications", "test_reports",
		"standards_applied", "declaration_of_conformity", "component_list",
		"classification_report", "mitigation_report", "verification_report",
		"evidence_index", "instructions_for_use", "monitoring_plan",
		"ai_intended_purpose", "ai_model_description", "ai_risk_management",
		"ai_human_oversight",
	}
	for _, st := range knownTypes {
		t.Run(st, func(t *testing.T) {
			prompt := getSystemPrompt(st)
			if prompt == "" {
				t.Errorf("getSystemPrompt(%q) returned empty string", st)
			}
		})
	}
}

func TestBuildUserPrompt_AIIntendedPurpose(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "ai_intended_purpose")
	if prompt == "" {
		t.Error("prompt should be non-empty for ai_intended_purpose")
	}
	if !strings.Contains(prompt, "Robot Arm XY-200") {
		t.Error("AI intended purpose prompt should contain machine name")
	}
}

func TestBuildUserPrompt_AIHumanOversight(t *testing.T) {
	sctx := newTestSectionContext()
	prompt := buildUserPrompt(sctx, "ai_human_oversight")
	if prompt == "" {
		t.Error("prompt should be non-empty for ai_human_oversight")
	}
	if !strings.Contains(prompt, "Robot Arm XY-200") {
		t.Error("AI human oversight prompt should contain machine name")
	}
}

func TestBuildUserPrompt_MultipleHazardAssessments(t *testing.T) {
	sctx := newTestSectionContext()

	// Find the high-risk hazard ID from the assessments map
	var highHazardID uuid.UUID
	for hid, assessments := range sctx.Assessments {
		if len(assessments) > 0 && assessments[0].RiskLevel == RiskLevelHigh {
			highHazardID = hid
			break
		}
	}

	if highHazardID != uuid.Nil {
		// Add a second assessment version (post-mitigation) for the high-risk hazard
		sctx.Assessments[highHazardID] = append(sctx.Assessments[highHazardID], RiskAssessment{
			ID:             uuid.New(),
			HazardID:       highHazardID,
			Version:        2,
			AssessmentType: AssessmentTypePostMitigation,
			Severity:       5,
			Exposure:       4,
			Probability:    1,
			Avoidance:      4,
			InherentRisk:   120.0,
			ResidualRisk:   20.0,
			RiskLevel:      RiskLevelMedium,
			IsAcceptable:   true,
		})
	}

	prompt := buildUserPrompt(sctx, "risk_assessment_report")
	if prompt == "" {
		t.Error("prompt should not be empty with multiple assessments")
	}
}