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

package iace

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

// ============================================================================
// Hazard CRUD Operations
// ============================================================================

// CreateHazard creates a new hazard within a project
func (s *Store) CreateHazard(ctx context.Context, req CreateHazardRequest) (*Hazard, error) {
	h := &Hazard{
		ID:              uuid.New(),
		ProjectID:       req.ProjectID,
		ComponentID:     req.ComponentID,
		LibraryHazardID: req.LibraryHazardID,
		Name:            req.Name,
		Description:     req.Description,
		Scenario:        req.Scenario,
		Category:        req.Category,
		SubCategory:     req.SubCategory,
		Status:          HazardStatusIdentified,
		MachineModule:   req.MachineModule,
		Function:        req.Function,
		LifecyclePhase:  req.LifecyclePhase,
		HazardousZone:   req.HazardousZone,
		TriggerEvent:    req.TriggerEvent,
		AffectedPerson:  req.AffectedPerson,
		PossibleHarm:    req.PossibleHarm,
		ReviewStatus:    ReviewStatusDraft,
		CreatedAt:       time.Now().UTC(),
		UpdatedAt:       time.Now().UTC(),
	}

	_, err := s.pool.Exec(ctx, `
		INSERT INTO iace_hazards (
			id, project_id, component_id, library_hazard_id,
			name, description, scenario, category, sub_category, status,
			machine_module, function, lifecycle_phase, hazardous_zone,
			trigger_event, affected_person, possible_harm, review_status,
			created_at, updated_at
		) VALUES (
			$1, $2, $3, $4,
			$5, $6, $7, $8, $9, $10,
			$11, $12, $13, $14,
			$15, $16, $17, $18,
			$19, $20
		)
	`,
		h.ID, h.ProjectID, h.ComponentID, h.LibraryHazardID,
		h.Name, h.Description, h.Scenario, h.Category, h.SubCategory, string(h.Status),
		h.MachineModule, h.Function, h.LifecyclePhase, h.HazardousZone,
		h.TriggerEvent, h.AffectedPerson, h.PossibleHarm, string(h.ReviewStatus),
		h.CreatedAt, h.UpdatedAt,
	)
	if err != nil {
		return nil, fmt.Errorf("create hazard: %w", err)
	}

	return h, nil
}

// GetHazard retrieves a hazard by ID
func (s *Store) GetHazard(ctx context.Context, id uuid.UUID) (*Hazard, error) {
	var h Hazard
	var status, reviewStatus string

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, project_id, component_id, library_hazard_id,
			name, description, scenario, category, sub_category, status,
			machine_module, function, lifecycle_phase, hazardous_zone,
			trigger_event, affected_person, possible_harm, review_status,
			created_at, updated_at
		FROM iace_hazards WHERE id = $1
	`, id).Scan(
		&h.ID, &h.ProjectID, &h.ComponentID, &h.LibraryHazardID,
		&h.Name, &h.Description, &h.Scenario, &h.Category, &h.SubCategory, &status,
		&h.MachineModule, &h.Function, &h.LifecyclePhase, &h.HazardousZone,
		&h.TriggerEvent, &h.AffectedPerson, &h.PossibleHarm, &reviewStatus,
		&h.CreatedAt, &h.UpdatedAt,
	)
	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("get hazard: %w", err)
	}

	h.Status = HazardStatus(status)
	h.ReviewStatus = ReviewStatus(reviewStatus)
	return &h, nil
}

// ListHazards lists all hazards for a project
func (s *Store) ListHazards(ctx context.Context, projectID uuid.UUID) ([]Hazard, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT
			id, project_id, component_id, library_hazard_id,
			name, description, scenario, category, sub_category, status,
			machine_module, function, lifecycle_phase, hazardous_zone,
			trigger_event, affected_person, possible_harm, review_status,
			created_at, updated_at
		FROM iace_hazards WHERE project_id = $1
		ORDER BY created_at ASC
	`, projectID)
	if err != nil {
		return nil, fmt.Errorf("list hazards: %w", err)
	}
	defer rows.Close()

	var hazards []Hazard
	for rows.Next() {
		var h Hazard
		var status, reviewStatus string

		err := rows.Scan(
			&h.ID, &h.ProjectID, &h.ComponentID, &h.LibraryHazardID,
			&h.Name, &h.Description, &h.Scenario, &h.Category, &h.SubCategory, &status,
			&h.MachineModule, &h.Function, &h.LifecyclePhase, &h.HazardousZone,
			&h.TriggerEvent, &h.AffectedPerson, &h.PossibleHarm, &reviewStatus,
			&h.CreatedAt, &h.UpdatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("list hazards scan: %w", err)
		}

		h.Status = HazardStatus(status)
		h.ReviewStatus = ReviewStatus(reviewStatus)
		hazards = append(hazards, h)
	}

	return hazards, nil
}

// UpdateHazard updates a hazard with a dynamic set of fields
func (s *Store) UpdateHazard(ctx context.Context, id uuid.UUID, updates map[string]interface{}) (*Hazard, error) {
	if len(updates) == 0 {
		return s.GetHazard(ctx, id)
	}

	query := "UPDATE iace_hazards SET updated_at = NOW()"
	args := []interface{}{id}
	argIdx := 2

	allowedFields := map[string]bool{
		"name": true, "description": true, "scenario": true, "category": true,
		"sub_category": true, "status": true, "component_id": true,
		"machine_module": true, "function": true, "lifecycle_phase": true,
		"hazardous_zone": true, "trigger_event": true, "affected_person": true,
		"possible_harm": true, "review_status": true,
	}

	for key, val := range updates {
		if allowedFields[key] {
			query += fmt.Sprintf(", %s = $%d", key, argIdx)
			args = append(args, val)
			argIdx++
		}
	}

	query += " WHERE id = $1"

	_, err := s.pool.Exec(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("update hazard: %w", err)
	}

	return s.GetHazard(ctx, id)
}

// ============================================================================
// Risk Assessment Operations
// ============================================================================

// CreateRiskAssessment creates a new risk assessment for a hazard
func (s *Store) CreateRiskAssessment(ctx context.Context, assessment *RiskAssessment) error {
	if assessment.ID == uuid.Nil {
		assessment.ID = uuid.New()
	}
	if assessment.CreatedAt.IsZero() {
		assessment.CreatedAt = time.Now().UTC()
	}

	_, err := s.pool.Exec(ctx, `
		INSERT INTO iace_risk_assessments (
			id, hazard_id, version, assessment_type,
			severity, exposure, probability,
			inherent_risk, control_maturity, control_coverage,
			test_evidence_strength, c_eff, residual_risk,
			risk_level, is_acceptable, acceptance_justification,
			assessed_by, created_at
		) VALUES (
			$1, $2, $3, $4,
			$5, $6, $7,
			$8, $9, $10,
			$11, $12, $13,
			$14, $15, $16,
			$17, $18
		)
	`,
		assessment.ID, assessment.HazardID, assessment.Version, string(assessment.AssessmentType),
		assessment.Severity, assessment.Exposure, assessment.Probability,
		assessment.InherentRisk, assessment.ControlMaturity, assessment.ControlCoverage,
		assessment.TestEvidenceStrength, assessment.CEff, assessment.ResidualRisk,
		string(assessment.RiskLevel), assessment.IsAcceptable, assessment.AcceptanceJustification,
		assessment.AssessedBy, assessment.CreatedAt,
	)
	if err != nil {
		return fmt.Errorf("create risk assessment: %w", err)
	}

	return nil
}

// GetLatestAssessment retrieves the most recent risk assessment for a hazard
func (s *Store) GetLatestAssessment(ctx context.Context, hazardID uuid.UUID) (*RiskAssessment, error) {
	var a RiskAssessment
	var assessmentType, riskLevel string

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, hazard_id, version, assessment_type,
			severity, exposure, probability,
			inherent_risk, control_maturity, control_coverage,
			test_evidence_strength, c_eff, residual_risk,
			risk_level, is_acceptable, acceptance_justification,
			assessed_by, created_at
		FROM iace_risk_assessments
		WHERE hazard_id = $1
		ORDER BY version DESC, created_at DESC
		LIMIT 1
	`, hazardID).Scan(
		&a.ID, &a.HazardID, &a.Version, &assessmentType,
		&a.Severity, &a.Exposure, &a.Probability,
		&a.InherentRisk, &a.ControlMaturity, &a.ControlCoverage,
		&a.TestEvidenceStrength, &a.CEff, &a.ResidualRisk,
		&riskLevel, &a.IsAcceptable, &a.AcceptanceJustification,
		&a.AssessedBy, &a.CreatedAt,
	)
	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("get latest assessment: %w", err)
	}

	a.AssessmentType = AssessmentType(assessmentType)
	a.RiskLevel = RiskLevel(riskLevel)

	return &a, nil
}

// ListAssessments lists all risk assessments for a hazard, newest first
func (s *Store) ListAssessments(ctx context.Context, hazardID uuid.UUID) ([]RiskAssessment, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT
			id, hazard_id, version, assessment_type,
			severity, exposure, probability,
			inherent_risk, control_maturity, control_coverage,
			test_evidence_strength, c_eff, residual_risk,
			risk_level, is_acceptable, acceptance_justification,
			assessed_by, created_at
		FROM iace_risk_assessments
		WHERE hazard_id = $1
		ORDER BY version DESC, created_at DESC
	`, hazardID)
	if err != nil {
		return nil, fmt.Errorf("list assessments: %w", err)
	}
	defer rows.Close()

	var assessments []RiskAssessment
	for rows.Next() {
		var a RiskAssessment
		var assessmentType, riskLevel string

		err := rows.Scan(
			&a.ID, &a.HazardID, &a.Version, &assessmentType,
			&a.Severity, &a.Exposure, &a.Probability,
			&a.InherentRisk, &a.ControlMaturity, &a.ControlCoverage,
			&a.TestEvidenceStrength, &a.CEff, &a.ResidualRisk,
			&riskLevel, &a.IsAcceptable, &a.AcceptanceJustification,
			&a.AssessedBy, &a.CreatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("list assessments scan: %w", err)
		}

		a.AssessmentType = AssessmentType(assessmentType)
		a.RiskLevel = RiskLevel(riskLevel)

		assessments = append(assessments, a)
	}

	return assessments, nil
}

// ============================================================================
// Risk Summary (Aggregated View)
// ============================================================================

// GetRiskSummary computes an aggregated risk overview for a project
func (s *Store) GetRiskSummary(ctx context.Context, projectID uuid.UUID) (*RiskSummaryResponse, error) {
	// Get all hazards for the project
	hazards, err := s.ListHazards(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("get risk summary - list hazards: %w", err)
	}

	summary := &RiskSummaryResponse{
		TotalHazards:  len(hazards),
		AllAcceptable: true,
	}

	if len(hazards) == 0 {
		summary.OverallRiskLevel = RiskLevelNegligible
		return summary, nil
	}

	highestRisk := RiskLevelNegligible

	for _, h := range hazards {
		latest, err := s.GetLatestAssessment(ctx, h.ID)
		if err != nil {
			return nil, fmt.Errorf("get risk summary - get assessment for hazard %s: %w", h.ID, err)
		}
		if latest == nil {
			// Hazard without assessment counts as unassessed; consider it not acceptable
			summary.AllAcceptable = false
			continue
		}

		switch latest.RiskLevel {
		case RiskLevelNotAcceptable:
			summary.NotAcceptable++
		case RiskLevelVeryHigh:
			summary.VeryHigh++
		case RiskLevelCritical:
			summary.Critical++
		case RiskLevelHigh:
			summary.High++
		case RiskLevelMedium:
			summary.Medium++
		case RiskLevelLow:
			summary.Low++
		case RiskLevelNegligible:
			summary.Negligible++
		}

		if !latest.IsAcceptable {
			summary.AllAcceptable = false
		}

		// Track highest risk level
		if riskLevelSeverity(latest.RiskLevel) > riskLevelSeverity(highestRisk) {
			highestRisk = latest.RiskLevel
		}
	}

	summary.OverallRiskLevel = highestRisk

	return summary, nil
}

// riskLevelSeverity returns a numeric severity for risk level comparison
func riskLevelSeverity(rl RiskLevel) int {
	switch rl {
	case RiskLevelNotAcceptable:
		return 7
	case RiskLevelVeryHigh:
		return 6
	case RiskLevelCritical:
		return 5
	case RiskLevelHigh:
		return 4
	case RiskLevelMedium:
		return 3
	case RiskLevelLow:
		return 2
	case RiskLevelNegligible:
		return 1
	default:
		return 0
	}
}

// ============================================================================
// Hazard Library Operations
// ============================================================================

// ListHazardLibrary lists hazard library entries, optionally filtered by category and component type
func (s *Store) ListHazardLibrary(ctx context.Context, category string, componentType string) ([]HazardLibraryEntry, error) {
	query := `
		SELECT
			id, category, COALESCE(sub_category, ''), name, description,
			default_severity, default_probability,
			COALESCE(default_exposure, 3), COALESCE(default_avoidance, 3),
			applicable_component_types, regulation_references,
			suggested_mitigations,
			COALESCE(typical_causes, '[]'::jsonb),
			COALESCE(typical_harm, ''),
			COALESCE(relevant_lifecycle_phases, '[]'::jsonb),
			COALESCE(recommended_measures_design, '[]'::jsonb),
			COALESCE(recommended_measures_technical, '[]'::jsonb),
			COALESCE(recommended_measures_information, '[]'::jsonb),
			COALESCE(suggested_evidence, '[]'::jsonb),
			COALESCE(related_keywords, '[]'::jsonb),
			is_builtin, tenant_id,
			created_at
		FROM iace_hazard_library WHERE 1=1`

	args := []interface{}{}
	argIdx := 1

	if category != "" {
		query += fmt.Sprintf(" AND category = $%d", argIdx)
		args = append(args, category)
		argIdx++
	}
	if componentType != "" {
		query += fmt.Sprintf(" AND applicable_component_types @> $%d::jsonb", argIdx)
		componentTypeJSON, _ := json.Marshal([]string{componentType})
		args = append(args, string(componentTypeJSON))
		argIdx++
	}

	query += " ORDER BY category ASC, name ASC"

	rows, err := s.pool.Query(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("list hazard library: %w", err)
	}
	defer rows.Close()

	var entries []HazardLibraryEntry
	for rows.Next() {
		var e HazardLibraryEntry
		var applicableComponentTypes, regulationReferences, suggestedMitigations []byte
		var typicalCauses, relevantPhases, measuresDesign, measuresTechnical, measuresInfo, evidence, keywords []byte

		err := rows.Scan(
			&e.ID, &e.Category, &e.SubCategory, &e.Name, &e.Description,
			&e.DefaultSeverity, &e.DefaultProbability,
			&e.DefaultExposure, &e.DefaultAvoidance,
			&applicableComponentTypes, &regulationReferences,
			&suggestedMitigations,
			&typicalCauses, &e.TypicalHarm, &relevantPhases,
			&measuresDesign, &measuresTechnical, &measuresInfo,
			&evidence, &keywords,
			&e.IsBuiltin, &e.TenantID,
			&e.CreatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("list hazard library scan: %w", err)
		}

		json.Unmarshal(applicableComponentTypes, &e.ApplicableComponentTypes)
		json.Unmarshal(regulationReferences, &e.RegulationReferences)
		json.Unmarshal(suggestedMitigations, &e.SuggestedMitigations)
		json.Unmarshal(typicalCauses, &e.TypicalCauses)
		json.Unmarshal(relevantPhases, &e.RelevantLifecyclePhases)
		json.Unmarshal(measuresDesign, &e.RecommendedMeasuresDesign)
		json.Unmarshal(measuresTechnical, &e.RecommendedMeasuresTechnical)
		json.Unmarshal(measuresInfo, &e.RecommendedMeasuresInformation)
		json.Unmarshal(evidence, &e.SuggestedEvidence)
		json.Unmarshal(keywords, &e.RelatedKeywords)

		if e.ApplicableComponentTypes == nil {
			e.ApplicableComponentTypes = []string{}
		}
		if e.RegulationReferences == nil {
			e.RegulationReferences = []string{}
		}

		entries = append(entries, e)
	}

	return entries, nil
}

// GetHazardLibraryEntry retrieves a single hazard library entry by ID
func (s *Store) GetHazardLibraryEntry(ctx context.Context, id uuid.UUID) (*HazardLibraryEntry, error) {
	var e HazardLibraryEntry
	var applicableComponentTypes, regulationReferences, suggestedMitigations []byte
	var typicalCauses, relevantLifecyclePhases []byte
	var recommendedMeasuresDesign, recommendedMeasuresTechnical, recommendedMeasuresInformation []byte
	var suggestedEvidence, relatedKeywords []byte

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, category, name, description,
			default_severity, default_probability,
			applicable_component_types, regulation_references,
			suggested_mitigations, is_builtin, tenant_id,
			created_at,
			COALESCE(sub_category, ''),
			COALESCE(default_exposure, 3),
			COALESCE(default_avoidance, 3),
			COALESCE(typical_causes, '[]'),
			COALESCE(typical_harm, ''),
			COALESCE(relevant_lifecycle_phases, '[]'),
			COALESCE(recommended_measures_design, '[]'),
			COALESCE(recommended_measures_technical, '[]'),
			COALESCE(recommended_measures_information, '[]'),
			COALESCE(suggested_evidence, '[]'),
			COALESCE(related_keywords, '[]')
		FROM iace_hazard_library WHERE id = $1
	`, id).Scan(
		&e.ID, &e.Category, &e.Name, &e.Description,
		&e.DefaultSeverity, &e.DefaultProbability,
		&applicableComponentTypes, &regulationReferences,
		&suggestedMitigations, &e.IsBuiltin, &e.TenantID,
		&e.CreatedAt,
		&e.SubCategory,
		&e.DefaultExposure, &e.DefaultAvoidance,
		&typicalCauses, &e.TypicalHarm,
		&relevantLifecyclePhases,
		&recommendedMeasuresDesign, &recommendedMeasuresTechnical, &recommendedMeasuresInformation,
		&suggestedEvidence, &relatedKeywords,
	)
	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("get hazard library entry: %w", err)
	}

	json.Unmarshal(applicableComponentTypes, &e.ApplicableComponentTypes)
	json.Unmarshal(regulationReferences, &e.RegulationReferences)
	json.Unmarshal(suggestedMitigations, &e.SuggestedMitigations)
	json.Unmarshal(typicalCauses, &e.TypicalCauses)
	json.Unmarshal(relevantLifecyclePhases, &e.RelevantLifecyclePhases)
	json.Unmarshal(recommendedMeasuresDesign, &e.RecommendedMeasuresDesign)
	json.Unmarshal(recommendedMeasuresTechnical, &e.RecommendedMeasuresTechnical)
	json.Unmarshal(recommendedMeasuresInformation, &e.RecommendedMeasuresInformation)
	json.Unmarshal(suggestedEvidence, &e.SuggestedEvidence)
	json.Unmarshal(relatedKeywords, &e.RelatedKeywords)

	if e.ApplicableComponentTypes == nil {
		e.ApplicableComponentTypes = []string{}
	}
	if e.RegulationReferences == nil {
		e.RegulationReferences = []string{}
	}

	return &e, nil
}