breakpilot-compliance/ai-compliance-sdk/internal/ucca/registration_store.go

package ucca

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

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

// AIRegistration represents an EU AI Database registration entry
type AIRegistration struct {
	ID        uuid.UUID `json:"id"`
	TenantID  uuid.UUID `json:"tenant_id"`

	// System
	SystemName        string `json:"system_name"`
	SystemVersion     string `json:"system_version,omitempty"`
	SystemDescription string `json:"system_description,omitempty"`
	IntendedPurpose   string `json:"intended_purpose,omitempty"`

	// Provider
	ProviderName            string `json:"provider_name,omitempty"`
	ProviderLegalForm       string `json:"provider_legal_form,omitempty"`
	ProviderAddress         string `json:"provider_address,omitempty"`
	ProviderCountry         string `json:"provider_country,omitempty"`
	EURepresentativeName    string `json:"eu_representative_name,omitempty"`
	EURepresentativeContact string `json:"eu_representative_contact,omitempty"`

	// Classification
	RiskClassification string `json:"risk_classification"`
	AnnexIIICategory   string `json:"annex_iii_category,omitempty"`
	GPAIClassification string `json:"gpai_classification"`

	// Conformity
	ConformityAssessmentType string `json:"conformity_assessment_type,omitempty"`
	NotifiedBodyName         string `json:"notified_body_name,omitempty"`
	NotifiedBodyID           string `json:"notified_body_id,omitempty"`
	CEMarking                bool   `json:"ce_marking"`

	// Training data
	TrainingDataCategories json.RawMessage `json:"training_data_categories,omitempty"`
	TrainingDataSummary    string          `json:"training_data_summary,omitempty"`

	// Status
	RegistrationStatus string     `json:"registration_status"`
	EUDatabaseID       string     `json:"eu_database_id,omitempty"`
	RegistrationDate   *time.Time `json:"registration_date,omitempty"`
	LastUpdateDate     *time.Time `json:"last_update_date,omitempty"`

	// Links
	UCCAAssessmentID     *uuid.UUID `json:"ucca_assessment_id,omitempty"`
	DecisionTreeResultID *uuid.UUID `json:"decision_tree_result_id,omitempty"`

	// Export
	ExportData json.RawMessage `json:"export_data,omitempty"`

	// Audit
	CreatedAt   time.Time `json:"created_at"`
	UpdatedAt   time.Time `json:"updated_at"`
	CreatedBy   string    `json:"created_by,omitempty"`
	SubmittedBy string    `json:"submitted_by,omitempty"`
}

// RegistrationStore handles AI registration persistence
type RegistrationStore struct {
	pool *pgxpool.Pool
}

// NewRegistrationStore creates a new registration store
func NewRegistrationStore(pool *pgxpool.Pool) *RegistrationStore {
	return &RegistrationStore{pool: pool}
}

// Create creates a new registration
func (s *RegistrationStore) Create(ctx context.Context, r *AIRegistration) error {
	r.ID = uuid.New()
	r.CreatedAt = time.Now()
	r.UpdatedAt = time.Now()
	if r.RegistrationStatus == "" {
		r.RegistrationStatus = "draft"
	}
	if r.RiskClassification == "" {
		r.RiskClassification = "not_classified"
	}
	if r.GPAIClassification == "" {
		r.GPAIClassification = "none"
	}

	_, err := s.pool.Exec(ctx, `
		INSERT INTO ai_system_registrations (
			id, tenant_id, system_name, system_version, system_description, intended_purpose,
			provider_name, provider_legal_form, provider_address, provider_country,
			eu_representative_name, eu_representative_contact,
			risk_classification, annex_iii_category, gpai_classification,
			conformity_assessment_type, notified_body_name, notified_body_id, ce_marking,
			training_data_categories, training_data_summary,
			registration_status, ucca_assessment_id, decision_tree_result_id,
			created_by
		) VALUES (
			$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12,
			$13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25
		)`,
		r.ID, r.TenantID, r.SystemName, r.SystemVersion, r.SystemDescription, r.IntendedPurpose,
		r.ProviderName, r.ProviderLegalForm, r.ProviderAddress, r.ProviderCountry,
		r.EURepresentativeName, r.EURepresentativeContact,
		r.RiskClassification, r.AnnexIIICategory, r.GPAIClassification,
		r.ConformityAssessmentType, r.NotifiedBodyName, r.NotifiedBodyID, r.CEMarking,
		r.TrainingDataCategories, r.TrainingDataSummary,
		r.RegistrationStatus, r.UCCAAssessmentID, r.DecisionTreeResultID,
		r.CreatedBy,
	)
	return err
}

// List returns all registrations for a tenant
func (s *RegistrationStore) List(ctx context.Context, tenantID uuid.UUID) ([]AIRegistration, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, tenant_id, system_name, system_version, system_description, intended_purpose,
			provider_name, provider_legal_form, provider_address, provider_country,
			eu_representative_name, eu_representative_contact,
			risk_classification, annex_iii_category, gpai_classification,
			conformity_assessment_type, notified_body_name, notified_body_id, ce_marking,
			training_data_categories, training_data_summary,
			registration_status, eu_database_id, registration_date, last_update_date,
			ucca_assessment_id, decision_tree_result_id, export_data,
			created_at, updated_at, created_by, submitted_by
		FROM ai_system_registrations
		WHERE tenant_id = $1
		ORDER BY created_at DESC`,
		tenantID,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var registrations []AIRegistration
	for rows.Next() {
		var r AIRegistration
		err := rows.Scan(
			&r.ID, &r.TenantID, &r.SystemName, &r.SystemVersion, &r.SystemDescription, &r.IntendedPurpose,
			&r.ProviderName, &r.ProviderLegalForm, &r.ProviderAddress, &r.ProviderCountry,
			&r.EURepresentativeName, &r.EURepresentativeContact,
			&r.RiskClassification, &r.AnnexIIICategory, &r.GPAIClassification,
			&r.ConformityAssessmentType, &r.NotifiedBodyName, &r.NotifiedBodyID, &r.CEMarking,
			&r.TrainingDataCategories, &r.TrainingDataSummary,
			&r.RegistrationStatus, &r.EUDatabaseID, &r.RegistrationDate, &r.LastUpdateDate,
			&r.UCCAAssessmentID, &r.DecisionTreeResultID, &r.ExportData,
			&r.CreatedAt, &r.UpdatedAt, &r.CreatedBy, &r.SubmittedBy,
		)
		if err != nil {
			return nil, err
		}
		registrations = append(registrations, r)
	}
	return registrations, nil
}

// GetByID returns a registration by ID
func (s *RegistrationStore) GetByID(ctx context.Context, id uuid.UUID) (*AIRegistration, error) {
	var r AIRegistration
	err := s.pool.QueryRow(ctx, `
		SELECT id, tenant_id, system_name, system_version, system_description, intended_purpose,
			provider_name, provider_legal_form, provider_address, provider_country,
			eu_representative_name, eu_representative_contact,
			risk_classification, annex_iii_category, gpai_classification,
			conformity_assessment_type, notified_body_name, notified_body_id, ce_marking,
			training_data_categories, training_data_summary,
			registration_status, eu_database_id, registration_date, last_update_date,
			ucca_assessment_id, decision_tree_result_id, export_data,
			created_at, updated_at, created_by, submitted_by
		FROM ai_system_registrations
		WHERE id = $1`,
		id,
	).Scan(
		&r.ID, &r.TenantID, &r.SystemName, &r.SystemVersion, &r.SystemDescription, &r.IntendedPurpose,
		&r.ProviderName, &r.ProviderLegalForm, &r.ProviderAddress, &r.ProviderCountry,
		&r.EURepresentativeName, &r.EURepresentativeContact,
		&r.RiskClassification, &r.AnnexIIICategory, &r.GPAIClassification,
		&r.ConformityAssessmentType, &r.NotifiedBodyName, &r.NotifiedBodyID, &r.CEMarking,
		&r.TrainingDataCategories, &r.TrainingDataSummary,
		&r.RegistrationStatus, &r.EUDatabaseID, &r.RegistrationDate, &r.LastUpdateDate,
		&r.UCCAAssessmentID, &r.DecisionTreeResultID, &r.ExportData,
		&r.CreatedAt, &r.UpdatedAt, &r.CreatedBy, &r.SubmittedBy,
	)
	if err != nil {
		return nil, err
	}
	return &r, nil
}

// Update updates a registration
func (s *RegistrationStore) Update(ctx context.Context, r *AIRegistration) error {
	r.UpdatedAt = time.Now()
	_, err := s.pool.Exec(ctx, `
		UPDATE ai_system_registrations SET
			system_name = $2, system_version = $3, system_description = $4, intended_purpose = $5,
			provider_name = $6, provider_legal_form = $7, provider_address = $8, provider_country = $9,
			eu_representative_name = $10, eu_representative_contact = $11,
			risk_classification = $12, annex_iii_category = $13, gpai_classification = $14,
			conformity_assessment_type = $15, notified_body_name = $16, notified_body_id = $17, ce_marking = $18,
			training_data_categories = $19, training_data_summary = $20,
			registration_status = $21, eu_database_id = $22,
			export_data = $23, updated_at = $24, submitted_by = $25
		WHERE id = $1`,
		r.ID, r.SystemName, r.SystemVersion, r.SystemDescription, r.IntendedPurpose,
		r.ProviderName, r.ProviderLegalForm, r.ProviderAddress, r.ProviderCountry,
		r.EURepresentativeName, r.EURepresentativeContact,
		r.RiskClassification, r.AnnexIIICategory, r.GPAIClassification,
		r.ConformityAssessmentType, r.NotifiedBodyName, r.NotifiedBodyID, r.CEMarking,
		r.TrainingDataCategories, r.TrainingDataSummary,
		r.RegistrationStatus, r.EUDatabaseID,
		r.ExportData, r.UpdatedAt, r.SubmittedBy,
	)
	return err
}

// UpdateStatus changes only the registration status
func (s *RegistrationStore) UpdateStatus(ctx context.Context, id uuid.UUID, status string, submittedBy string) error {
	now := time.Now()
	_, err := s.pool.Exec(ctx, `
		UPDATE ai_system_registrations
		SET registration_status = $2, submitted_by = $3, updated_at = $4,
			registration_date = CASE WHEN $2 = 'submitted' THEN $4 ELSE registration_date END,
			last_update_date = $4
		WHERE id = $1`,
		id, status, submittedBy, now,
	)
	return err
}

// BuildExportJSON creates the EU AI Database submission JSON
func (s *RegistrationStore) BuildExportJSON(r *AIRegistration) json.RawMessage {
	export := map[string]interface{}{
		"schema_version": "1.0",
		"submission_type": "ai_system_registration",
		"regulation": "EU AI Act (EU) 2024/1689",
		"article": "Art. 49",
		"provider": map[string]interface{}{
			"name":               r.ProviderName,
			"legal_form":         r.ProviderLegalForm,
			"address":            r.ProviderAddress,
			"country":            r.ProviderCountry,
			"eu_representative":  r.EURepresentativeName,
			"eu_rep_contact":     r.EURepresentativeContact,
		},
		"system": map[string]interface{}{
			"name":        r.SystemName,
			"version":     r.SystemVersion,
			"description": r.SystemDescription,
			"purpose":     r.IntendedPurpose,
		},
		"classification": map[string]interface{}{
			"risk_level":         r.RiskClassification,
			"annex_iii_category": r.AnnexIIICategory,
			"gpai":               r.GPAIClassification,
		},
		"conformity": map[string]interface{}{
			"assessment_type": r.ConformityAssessmentType,
			"notified_body":   r.NotifiedBodyName,
			"notified_body_id": r.NotifiedBodyID,
			"ce_marking":      r.CEMarking,
		},
		"training_data": map[string]interface{}{
			"categories": r.TrainingDataCategories,
			"summary":    r.TrainingDataSummary,
		},
		"status": r.RegistrationStatus,
	}
	data, _ := json.Marshal(export)
	return data
}