breakpilot-pwa/ai-compliance-sdk/internal/incidents/store.go

package incidents

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

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

// Store handles incident data persistence
type Store struct {
	pool *pgxpool.Pool
}

// NewStore creates a new incident store
func NewStore(pool *pgxpool.Pool) *Store {
	return &Store{pool: pool}
}

// ============================================================================
// Incident CRUD Operations
// ============================================================================

// CreateIncident creates a new incident
func (s *Store) CreateIncident(ctx context.Context, incident *Incident) error {
	incident.ID = uuid.New()
	incident.CreatedAt = time.Now().UTC()
	incident.UpdatedAt = incident.CreatedAt
	if incident.Status == "" {
		incident.Status = IncidentStatusDetected
	}
	if incident.AffectedDataCategories == nil {
		incident.AffectedDataCategories = []string{}
	}
	if incident.AffectedSystems == nil {
		incident.AffectedSystems = []string{}
	}
	if incident.Timeline == nil {
		incident.Timeline = []TimelineEntry{}
	}

	affectedDataCategories, _ := json.Marshal(incident.AffectedDataCategories)
	affectedSystems, _ := json.Marshal(incident.AffectedSystems)
	riskAssessment, _ := json.Marshal(incident.RiskAssessment)
	authorityNotification, _ := json.Marshal(incident.AuthorityNotification)
	dataSubjectNotification, _ := json.Marshal(incident.DataSubjectNotification)
	timeline, _ := json.Marshal(incident.Timeline)

	_, err := s.pool.Exec(ctx, `
		INSERT INTO incident_incidents (
			id, tenant_id, title, description, category, status, severity,
			detected_at, reported_by,
			affected_data_categories, affected_data_subject_count, affected_systems,
			risk_assessment, authority_notification, data_subject_notification,
			root_cause, lessons_learned, timeline,
			created_at, updated_at, closed_at
		) VALUES (
			$1, $2, $3, $4, $5, $6, $7,
			$8, $9,
			$10, $11, $12,
			$13, $14, $15,
			$16, $17, $18,
			$19, $20, $21
		)
	`,
		incident.ID, incident.TenantID, incident.Title, incident.Description,
		string(incident.Category), string(incident.Status), string(incident.Severity),
		incident.DetectedAt, incident.ReportedBy,
		affectedDataCategories, incident.AffectedDataSubjectCount, affectedSystems,
		riskAssessment, authorityNotification, dataSubjectNotification,
		incident.RootCause, incident.LessonsLearned, timeline,
		incident.CreatedAt, incident.UpdatedAt, incident.ClosedAt,
	)

	return err
}

// GetIncident retrieves an incident by ID
func (s *Store) GetIncident(ctx context.Context, id uuid.UUID) (*Incident, error) {
	var incident Incident
	var category, status, severity string
	var affectedDataCategories, affectedSystems []byte
	var riskAssessment, authorityNotification, dataSubjectNotification []byte
	var timeline []byte

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, tenant_id, title, description, category, status, severity,
			detected_at, reported_by,
			affected_data_categories, affected_data_subject_count, affected_systems,
			risk_assessment, authority_notification, data_subject_notification,
			root_cause, lessons_learned, timeline,
			created_at, updated_at, closed_at
		FROM incident_incidents WHERE id = $1
	`, id).Scan(
		&incident.ID, &incident.TenantID, &incident.Title, &incident.Description,
		&category, &status, &severity,
		&incident.DetectedAt, &incident.ReportedBy,
		&affectedDataCategories, &incident.AffectedDataSubjectCount, &affectedSystems,
		&riskAssessment, &authorityNotification, &dataSubjectNotification,
		&incident.RootCause, &incident.LessonsLearned, &timeline,
		&incident.CreatedAt, &incident.UpdatedAt, &incident.ClosedAt,
	)

	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	incident.Category = IncidentCategory(category)
	incident.Status = IncidentStatus(status)
	incident.Severity = IncidentSeverity(severity)

	json.Unmarshal(affectedDataCategories, &incident.AffectedDataCategories)
	json.Unmarshal(affectedSystems, &incident.AffectedSystems)
	json.Unmarshal(riskAssessment, &incident.RiskAssessment)
	json.Unmarshal(authorityNotification, &incident.AuthorityNotification)
	json.Unmarshal(dataSubjectNotification, &incident.DataSubjectNotification)
	json.Unmarshal(timeline, &incident.Timeline)

	if incident.AffectedDataCategories == nil {
		incident.AffectedDataCategories = []string{}
	}
	if incident.AffectedSystems == nil {
		incident.AffectedSystems = []string{}
	}
	if incident.Timeline == nil {
		incident.Timeline = []TimelineEntry{}
	}

	return &incident, nil
}

// ListIncidents lists incidents for a tenant with optional filters
func (s *Store) ListIncidents(ctx context.Context, tenantID uuid.UUID, filters *IncidentFilters) ([]Incident, int, error) {
	// Count query
	countQuery := "SELECT COUNT(*) FROM incident_incidents WHERE tenant_id = $1"
	countArgs := []interface{}{tenantID}
	countArgIdx := 2

	if filters != nil {
		if filters.Status != "" {
			countQuery += fmt.Sprintf(" AND status = $%d", countArgIdx)
			countArgs = append(countArgs, string(filters.Status))
			countArgIdx++
		}
		if filters.Severity != "" {
			countQuery += fmt.Sprintf(" AND severity = $%d", countArgIdx)
			countArgs = append(countArgs, string(filters.Severity))
			countArgIdx++
		}
		if filters.Category != "" {
			countQuery += fmt.Sprintf(" AND category = $%d", countArgIdx)
			countArgs = append(countArgs, string(filters.Category))
			countArgIdx++
		}
	}

	var total int
	err := s.pool.QueryRow(ctx, countQuery, countArgs...).Scan(&total)
	if err != nil {
		return nil, 0, err
	}

	// Data query
	query := `
		SELECT
			id, tenant_id, title, description, category, status, severity,
			detected_at, reported_by,
			affected_data_categories, affected_data_subject_count, affected_systems,
			risk_assessment, authority_notification, data_subject_notification,
			root_cause, lessons_learned, timeline,
			created_at, updated_at, closed_at
		FROM incident_incidents WHERE tenant_id = $1`

	args := []interface{}{tenantID}
	argIdx := 2

	if filters != nil {
		if filters.Status != "" {
			query += fmt.Sprintf(" AND status = $%d", argIdx)
			args = append(args, string(filters.Status))
			argIdx++
		}
		if filters.Severity != "" {
			query += fmt.Sprintf(" AND severity = $%d", argIdx)
			args = append(args, string(filters.Severity))
			argIdx++
		}
		if filters.Category != "" {
			query += fmt.Sprintf(" AND category = $%d", argIdx)
			args = append(args, string(filters.Category))
			argIdx++
		}
	}

	query += " ORDER BY detected_at DESC"

	if filters != nil && filters.Limit > 0 {
		query += fmt.Sprintf(" LIMIT $%d", argIdx)
		args = append(args, filters.Limit)
		argIdx++

		if filters.Offset > 0 {
			query += fmt.Sprintf(" OFFSET $%d", argIdx)
			args = append(args, filters.Offset)
			argIdx++
		}
	}

	rows, err := s.pool.Query(ctx, query, args...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	var incidents []Incident
	for rows.Next() {
		var incident Incident
		var category, status, severity string
		var affectedDataCategories, affectedSystems []byte
		var riskAssessment, authorityNotification, dataSubjectNotification []byte
		var timeline []byte

		err := rows.Scan(
			&incident.ID, &incident.TenantID, &incident.Title, &incident.Description,
			&category, &status, &severity,
			&incident.DetectedAt, &incident.ReportedBy,
			&affectedDataCategories, &incident.AffectedDataSubjectCount, &affectedSystems,
			&riskAssessment, &authorityNotification, &dataSubjectNotification,
			&incident.RootCause, &incident.LessonsLearned, &timeline,
			&incident.CreatedAt, &incident.UpdatedAt, &incident.ClosedAt,
		)
		if err != nil {
			return nil, 0, err
		}

		incident.Category = IncidentCategory(category)
		incident.Status = IncidentStatus(status)
		incident.Severity = IncidentSeverity(severity)

		json.Unmarshal(affectedDataCategories, &incident.AffectedDataCategories)
		json.Unmarshal(affectedSystems, &incident.AffectedSystems)
		json.Unmarshal(riskAssessment, &incident.RiskAssessment)
		json.Unmarshal(authorityNotification, &incident.AuthorityNotification)
		json.Unmarshal(dataSubjectNotification, &incident.DataSubjectNotification)
		json.Unmarshal(timeline, &incident.Timeline)

		if incident.AffectedDataCategories == nil {
			incident.AffectedDataCategories = []string{}
		}
		if incident.AffectedSystems == nil {
			incident.AffectedSystems = []string{}
		}
		if incident.Timeline == nil {
			incident.Timeline = []TimelineEntry{}
		}

		incidents = append(incidents, incident)
	}

	return incidents, total, nil
}

// UpdateIncident updates an incident
func (s *Store) UpdateIncident(ctx context.Context, incident *Incident) error {
	incident.UpdatedAt = time.Now().UTC()

	affectedDataCategories, _ := json.Marshal(incident.AffectedDataCategories)
	affectedSystems, _ := json.Marshal(incident.AffectedSystems)

	_, err := s.pool.Exec(ctx, `
		UPDATE incident_incidents SET
			title = $2, description = $3, category = $4, status = $5, severity = $6,
			affected_data_categories = $7, affected_data_subject_count = $8, affected_systems = $9,
			root_cause = $10, lessons_learned = $11,
			updated_at = $12
		WHERE id = $1
	`,
		incident.ID, incident.Title, incident.Description,
		string(incident.Category), string(incident.Status), string(incident.Severity),
		affectedDataCategories, incident.AffectedDataSubjectCount, affectedSystems,
		incident.RootCause, incident.LessonsLearned,
		incident.UpdatedAt,
	)

	return err
}

// DeleteIncident deletes an incident and its related measures (cascade handled by FK)
func (s *Store) DeleteIncident(ctx context.Context, id uuid.UUID) error {
	_, err := s.pool.Exec(ctx, "DELETE FROM incident_incidents WHERE id = $1", id)
	return err
}

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

// UpdateRiskAssessment updates the risk assessment for an incident
func (s *Store) UpdateRiskAssessment(ctx context.Context, incidentID uuid.UUID, assessment *RiskAssessment) error {
	assessmentJSON, _ := json.Marshal(assessment)

	_, err := s.pool.Exec(ctx, `
		UPDATE incident_incidents SET
			risk_assessment = $2,
			updated_at = NOW()
		WHERE id = $1
	`, incidentID, assessmentJSON)

	return err
}

// ============================================================================
// Notification Operations
// ============================================================================

// UpdateAuthorityNotification updates the authority notification for an incident
func (s *Store) UpdateAuthorityNotification(ctx context.Context, incidentID uuid.UUID, notification *AuthorityNotification) error {
	notificationJSON, _ := json.Marshal(notification)

	_, err := s.pool.Exec(ctx, `
		UPDATE incident_incidents SET
			authority_notification = $2,
			updated_at = NOW()
		WHERE id = $1
	`, incidentID, notificationJSON)

	return err
}

// UpdateDataSubjectNotification updates the data subject notification for an incident
func (s *Store) UpdateDataSubjectNotification(ctx context.Context, incidentID uuid.UUID, notification *DataSubjectNotification) error {
	notificationJSON, _ := json.Marshal(notification)

	_, err := s.pool.Exec(ctx, `
		UPDATE incident_incidents SET
			data_subject_notification = $2,
			updated_at = NOW()
		WHERE id = $1
	`, incidentID, notificationJSON)

	return err
}

// ============================================================================
// Measure Operations
// ============================================================================

// AddMeasure adds a corrective measure to an incident
func (s *Store) AddMeasure(ctx context.Context, measure *IncidentMeasure) error {
	measure.ID = uuid.New()
	measure.CreatedAt = time.Now().UTC()
	if measure.Status == "" {
		measure.Status = MeasureStatusPlanned
	}

	_, err := s.pool.Exec(ctx, `
		INSERT INTO incident_measures (
			id, incident_id, title, description, measure_type, status,
			responsible, due_date, completed_at, created_at
		) VALUES (
			$1, $2, $3, $4, $5, $6,
			$7, $8, $9, $10
		)
	`,
		measure.ID, measure.IncidentID, measure.Title, measure.Description,
		string(measure.MeasureType), string(measure.Status),
		measure.Responsible, measure.DueDate, measure.CompletedAt, measure.CreatedAt,
	)

	return err
}

// ListMeasures lists all measures for an incident
func (s *Store) ListMeasures(ctx context.Context, incidentID uuid.UUID) ([]IncidentMeasure, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT
			id, incident_id, title, description, measure_type, status,
			responsible, due_date, completed_at, created_at
		FROM incident_measures WHERE incident_id = $1
		ORDER BY created_at ASC
	`, incidentID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var measures []IncidentMeasure
	for rows.Next() {
		var m IncidentMeasure
		var measureType, status string

		err := rows.Scan(
			&m.ID, &m.IncidentID, &m.Title, &m.Description,
			&measureType, &status,
			&m.Responsible, &m.DueDate, &m.CompletedAt, &m.CreatedAt,
		)
		if err != nil {
			return nil, err
		}

		m.MeasureType = MeasureType(measureType)
		m.Status = MeasureStatus(status)

		measures = append(measures, m)
	}

	return measures, nil
}

// UpdateMeasure updates an existing measure
func (s *Store) UpdateMeasure(ctx context.Context, measure *IncidentMeasure) error {
	_, err := s.pool.Exec(ctx, `
		UPDATE incident_measures SET
			title = $2, description = $3, measure_type = $4, status = $5,
			responsible = $6, due_date = $7, completed_at = $8
		WHERE id = $1
	`,
		measure.ID, measure.Title, measure.Description,
		string(measure.MeasureType), string(measure.Status),
		measure.Responsible, measure.DueDate, measure.CompletedAt,
	)

	return err
}

// CompleteMeasure marks a measure as completed
func (s *Store) CompleteMeasure(ctx context.Context, id uuid.UUID) error {
	now := time.Now().UTC()

	_, err := s.pool.Exec(ctx, `
		UPDATE incident_measures SET
			status = $2,
			completed_at = $3
		WHERE id = $1
	`, id, string(MeasureStatusCompleted), now)

	return err
}

// ============================================================================
// Timeline Operations
// ============================================================================

// AddTimelineEntry appends a timeline entry to the incident's JSONB timeline array
func (s *Store) AddTimelineEntry(ctx context.Context, incidentID uuid.UUID, entry TimelineEntry) error {
	entryJSON, err := json.Marshal(entry)
	if err != nil {
		return err
	}

	// Use the || operator to append to the JSONB array
	_, err = s.pool.Exec(ctx, `
		UPDATE incident_incidents SET
			timeline = COALESCE(timeline, '[]'::jsonb) || $2::jsonb,
			updated_at = NOW()
		WHERE id = $1
	`, incidentID, string(entryJSON))

	return err
}

// ============================================================================
// Close Incident
// ============================================================================

// CloseIncident closes an incident with root cause and lessons learned
func (s *Store) CloseIncident(ctx context.Context, id uuid.UUID, rootCause, lessonsLearned string) error {
	now := time.Now().UTC()

	_, err := s.pool.Exec(ctx, `
		UPDATE incident_incidents SET
			status = $2,
			root_cause = $3,
			lessons_learned = $4,
			closed_at = $5,
			updated_at = $5
		WHERE id = $1
	`, id, string(IncidentStatusClosed), rootCause, lessonsLearned, now)

	return err
}

// ============================================================================
// Statistics
// ============================================================================

// GetStatistics returns aggregated incident statistics for a tenant
func (s *Store) GetStatistics(ctx context.Context, tenantID uuid.UUID) (*IncidentStatistics, error) {
	stats := &IncidentStatistics{
		ByStatus:   make(map[string]int),
		BySeverity: make(map[string]int),
		ByCategory: make(map[string]int),
	}

	// Total incidents
	s.pool.QueryRow(ctx,
		"SELECT COUNT(*) FROM incident_incidents WHERE tenant_id = $1",
		tenantID).Scan(&stats.TotalIncidents)

	// Open incidents (not closed)
	s.pool.QueryRow(ctx,
		"SELECT COUNT(*) FROM incident_incidents WHERE tenant_id = $1 AND status != 'closed'",
		tenantID).Scan(&stats.OpenIncidents)

	// By status
	rows, err := s.pool.Query(ctx,
		"SELECT status, COUNT(*) FROM incident_incidents WHERE tenant_id = $1 GROUP BY status",
		tenantID)
	if err == nil {
		defer rows.Close()
		for rows.Next() {
			var status string
			var count int
			rows.Scan(&status, &count)
			stats.ByStatus[status] = count
		}
	}

	// By severity
	rows, err = s.pool.Query(ctx,
		"SELECT severity, COUNT(*) FROM incident_incidents WHERE tenant_id = $1 GROUP BY severity",
		tenantID)
	if err == nil {
		defer rows.Close()
		for rows.Next() {
			var severity string
			var count int
			rows.Scan(&severity, &count)
			stats.BySeverity[severity] = count
		}
	}

	// By category
	rows, err = s.pool.Query(ctx,
		"SELECT category, COUNT(*) FROM incident_incidents WHERE tenant_id = $1 GROUP BY category",
		tenantID)
	if err == nil {
		defer rows.Close()
		for rows.Next() {
			var category string
			var count int
			rows.Scan(&category, &count)
			stats.ByCategory[category] = count
		}
	}

	// Notifications pending
	s.pool.QueryRow(ctx, `
		SELECT COUNT(*) FROM incident_incidents
		WHERE tenant_id = $1
		AND (authority_notification->>'status' = 'pending'
		     OR data_subject_notification->>'status' = 'pending')
	`, tenantID).Scan(&stats.NotificationsPending)

	// Average resolution hours (for closed incidents)
	s.pool.QueryRow(ctx, `
		SELECT COALESCE(AVG(EXTRACT(EPOCH FROM (closed_at - detected_at)) / 3600), 0)
		FROM incident_incidents
		WHERE tenant_id = $1 AND status = 'closed' AND closed_at IS NOT NULL
	`, tenantID).Scan(&stats.AvgResolutionHours)

	return stats, nil
}