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

package iace

import (
	"context"
	"fmt"
	"time"

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

// ============================================================================
// Mitigation CRUD Operations
// ============================================================================

// CreateMitigation creates a new mitigation measure for a hazard
func (s *Store) CreateMitigation(ctx context.Context, req CreateMitigationRequest) (*Mitigation, error) {
	m := &Mitigation{
		ID:            uuid.New(),
		HazardID:      req.HazardID,
		ReductionType: req.ReductionType,
		Name:          req.Name,
		Description:   req.Description,
		Status:        MitigationStatusPlanned,
		CreatedAt:     time.Now().UTC(),
		UpdatedAt:     time.Now().UTC(),
	}

	// ON CONFLICT DO NOTHING + RETURNING id makes the call idempotent against
	// the iace_mitigations_hazard_name_uniq constraint (migration 030).
	// A repeated Engine-init no longer creates duplicates; the second click
	// is silently swallowed. If nothing was inserted we look up the existing
	// row and return that, so the caller always gets a usable Mitigation.
	var insertedID uuid.UUID
	err := s.pool.QueryRow(ctx, `
		INSERT INTO iace_mitigations (
			id, hazard_id, reduction_type, name, description,
			status, verification_method, verification_result,
			verified_at, verified_by,
			is_relevant, is_customer_standard,
			created_at, updated_at
		) VALUES (
			$1, $2, $3, $4, $5,
			$6, $7, $8,
			$9, $10,
			$11, $12,
			$13, $14
		)
		ON CONFLICT (hazard_id, name) DO NOTHING
		RETURNING id
	`,
		m.ID, m.HazardID, string(m.ReductionType), m.Name, m.Description,
		string(m.Status), "", "",
		nil, uuid.Nil,
		m.IsRelevant, m.IsCustomerStandard,
		m.CreatedAt, m.UpdatedAt,
	).Scan(&insertedID)

	if err == pgx.ErrNoRows {
		// Conflict — return the existing row for the same (hazard_id, name).
		var existingID uuid.UUID
		if lookupErr := s.pool.QueryRow(ctx,
			`SELECT id FROM iace_mitigations WHERE hazard_id = $1 AND name = $2`,
			m.HazardID, m.Name,
		).Scan(&existingID); lookupErr != nil {
			return nil, fmt.Errorf("create mitigation lookup-on-conflict: %w", lookupErr)
		}
		return s.getMitigation(ctx, existingID)
	}
	if err != nil {
		return nil, fmt.Errorf("create mitigation: %w", err)
	}

	return m, nil
}

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

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

	for key, val := range updates {
		switch key {
		case "name", "description", "verification_result":
			query += fmt.Sprintf(", %s = $%d", key, argIdx)
			args = append(args, val)
			argIdx++
		case "status":
			query += fmt.Sprintf(", status = $%d", argIdx)
			args = append(args, val)
			argIdx++
		case "reduction_type":
			query += fmt.Sprintf(", reduction_type = $%d", argIdx)
			args = append(args, val)
			argIdx++
		case "verification_method":
			query += fmt.Sprintf(", verification_method = $%d", argIdx)
			args = append(args, val)
			argIdx++
		case "is_relevant":
			query += fmt.Sprintf(", is_relevant = $%d", argIdx)
			args = append(args, val)
			argIdx++
		case "is_customer_standard":
			// CHECK constraint requires is_relevant=true when this is true,
			// so we flip is_relevant on as well when the caller sets the
			// customer-standard flag.
			b, _ := val.(bool)
			query += fmt.Sprintf(", is_customer_standard = $%d", argIdx)
			args = append(args, b)
			argIdx++
			if b {
				query += fmt.Sprintf(", is_relevant = $%d", argIdx)
				args = append(args, true)
				argIdx++
			}
		}
	}

	query += " WHERE id = $1"

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

	return s.getMitigation(ctx, id)
}

// VerifyMitigation marks a mitigation as verified
func (s *Store) VerifyMitigation(ctx context.Context, id uuid.UUID, verificationResult string, verifiedBy string) error {
	now := time.Now().UTC()
	verifiedByUUID, err := uuid.Parse(verifiedBy)
	if err != nil {
		return fmt.Errorf("invalid verified_by UUID: %w", err)
	}

	_, err = s.pool.Exec(ctx, `
		UPDATE iace_mitigations SET
			status = $2,
			verification_result = $3,
			verified_at = $4,
			verified_by = $5,
			updated_at = $4
		WHERE id = $1
	`, id, string(MitigationStatusVerified), verificationResult, now, verifiedByUUID)
	if err != nil {
		return fmt.Errorf("verify mitigation: %w", err)
	}

	return nil
}

// ListMitigations lists all mitigations for a hazard
func (s *Store) ListMitigations(ctx context.Context, hazardID uuid.UUID) ([]Mitigation, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT
			id, hazard_id, reduction_type, name, description,
			status, verification_method, verification_result,
			verified_at, verified_by,
			is_relevant, is_customer_standard,
			created_at, updated_at
		FROM iace_mitigations WHERE hazard_id = $1
		ORDER BY created_at ASC
	`, hazardID)
	if err != nil {
		return nil, fmt.Errorf("list mitigations: %w", err)
	}
	defer rows.Close()

	var mitigations []Mitigation
	for rows.Next() {
		var m Mitigation
		var reductionType, status, verificationMethod string

		err := rows.Scan(
			&m.ID, &m.HazardID, &reductionType, &m.Name, &m.Description,
			&status, &verificationMethod, &m.VerificationResult,
			&m.VerifiedAt, &m.VerifiedBy,
			&m.IsRelevant, &m.IsCustomerStandard,
			&m.CreatedAt, &m.UpdatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("list mitigations scan: %w", err)
		}

		m.ReductionType = ReductionType(reductionType)
		m.Status = MitigationStatus(status)
		m.VerificationMethod = VerificationMethod(verificationMethod)

		mitigations = append(mitigations, m)
	}

	return mitigations, nil
}

// ListMitigationsByProject lists all mitigations for all hazards in a project.
func (s *Store) ListMitigationsByProject(ctx context.Context, projectID uuid.UUID) ([]Mitigation, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT
			m.id, m.hazard_id, m.reduction_type, m.name, m.description,
			m.status, m.verification_method, m.verification_result,
			m.verified_at, m.verified_by,
			m.is_relevant, m.is_customer_standard,
			m.created_at, m.updated_at
		FROM iace_mitigations m
		JOIN iace_hazards h ON h.id = m.hazard_id
		WHERE h.project_id = $1
		ORDER BY m.created_at ASC
	`, projectID)
	if err != nil {
		return nil, fmt.Errorf("list mitigations by project: %w", err)
	}
	defer rows.Close()

	var mitigations []Mitigation
	for rows.Next() {
		var m Mitigation
		var reductionType, status, verificationMethod string

		err := rows.Scan(
			&m.ID, &m.HazardID, &reductionType, &m.Name, &m.Description,
			&status, &verificationMethod, &m.VerificationResult,
			&m.VerifiedAt, &m.VerifiedBy,
			&m.IsRelevant, &m.IsCustomerStandard,
			&m.CreatedAt, &m.UpdatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("list mitigations by project scan: %w", err)
		}

		m.ReductionType = ReductionType(reductionType)
		m.Status = MitigationStatus(status)
		m.VerificationMethod = VerificationMethod(verificationMethod)

		mitigations = append(mitigations, m)
	}

	return mitigations, nil
}

// DeleteMitigation deletes a mitigation by ID.
func (s *Store) DeleteMitigation(ctx context.Context, id uuid.UUID) error {
	_, err := s.pool.Exec(ctx, `DELETE FROM iace_mitigations WHERE id = $1`, id)
	if err != nil {
		return fmt.Errorf("delete mitigation: %w", err)
	}
	return nil
}

// GetMitigation fetches a single mitigation by ID.
func (s *Store) GetMitigation(ctx context.Context, id uuid.UUID) (*Mitigation, error) {
	return s.getMitigation(ctx, id)
}

// getMitigation is a helper to fetch a single mitigation by ID
func (s *Store) getMitigation(ctx context.Context, id uuid.UUID) (*Mitigation, error) {
	var m Mitigation
	var reductionType, status, verificationMethod string

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, hazard_id, reduction_type, name, description,
			status, verification_method, verification_result,
			verified_at, verified_by,
			is_relevant, is_customer_standard,
			created_at, updated_at
		FROM iace_mitigations WHERE id = $1
	`, id).Scan(
		&m.ID, &m.HazardID, &reductionType, &m.Name, &m.Description,
		&status, &verificationMethod, &m.VerificationResult,
		&m.VerifiedAt, &m.VerifiedBy,
		&m.IsRelevant, &m.IsCustomerStandard,
		&m.CreatedAt, &m.UpdatedAt,
	)
	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("get mitigation: %w", err)
	}

	m.ReductionType = ReductionType(reductionType)
	m.Status = MitigationStatus(status)
	m.VerificationMethod = VerificationMethod(verificationMethod)

	return &m, nil
}