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

package ucca

import (
	"fmt"
	"sort"
	"time"

	"github.com/google/uuid"
)

// ============================================================================
// Obligations Registry
// ============================================================================
//
// The registry manages all regulation modules and provides methods to evaluate
// facts against all registered regulations, aggregating the results.
//
// ============================================================================

// ObligationsRegistry manages all regulation modules
type ObligationsRegistry struct {
	modules map[string]RegulationModule
}

// NewObligationsRegistry creates a new registry and registers all default modules.
// It loads v2 JSON modules first; for regulations without v2 JSON, falls back to YAML modules.
func NewObligationsRegistry() *ObligationsRegistry {
	r := &ObligationsRegistry{
		modules: make(map[string]RegulationModule),
	}

	// Try to load v2 JSON modules first
	v2Loaded := r.loadV2Modules()

	// Fall back to YAML modules for regulations not covered by v2
	if !v2Loaded["nis2"] {
		if nis2Module, err := NewNIS2Module(); err == nil {
			r.Register(nis2Module)
		} else {
			fmt.Printf("Warning: Could not load NIS2 module: %v\n", err)
		}
	}

	if !v2Loaded["dsgvo"] {
		if dsgvoModule, err := NewDSGVOModule(); err == nil {
			r.Register(dsgvoModule)
		} else {
			fmt.Printf("Warning: Could not load DSGVO module: %v\n", err)
		}
	}

	if !v2Loaded["ai_act"] {
		if aiActModule, err := NewAIActModule(); err == nil {
			r.Register(aiActModule)
		} else {
			fmt.Printf("Warning: Could not load AI Act module: %v\n", err)
		}
	}

	return r
}

// loadV2Modules attempts to load all v2 JSON regulation modules
func (r *ObligationsRegistry) loadV2Modules() map[string]bool {
	loaded := make(map[string]bool)

	regulations, err := LoadAllV2Regulations()
	if err != nil {
		fmt.Printf("Info: No v2 regulations found, using YAML modules: %v\n", err)
		return loaded
	}

	for regID, regFile := range regulations {
		module := NewJSONRegulationModule(regFile)
		r.Register(module)
		loaded[regID] = true
		fmt.Printf("Loaded v2 regulation module: %s (%d obligations)\n", regID, len(regFile.Obligations))
	}

	return loaded
}

// NewObligationsRegistryWithModules creates a registry with specific modules
func NewObligationsRegistryWithModules(modules ...RegulationModule) *ObligationsRegistry {
	r := &ObligationsRegistry{
		modules: make(map[string]RegulationModule),
	}
	for _, m := range modules {
		r.Register(m)
	}
	return r
}

// Register adds a regulation module to the registry
func (r *ObligationsRegistry) Register(module RegulationModule) {
	r.modules[module.ID()] = module
}

// Unregister removes a regulation module from the registry
func (r *ObligationsRegistry) Unregister(moduleID string) {
	delete(r.modules, moduleID)
}

// GetModule returns a specific module by ID
func (r *ObligationsRegistry) GetModule(moduleID string) (RegulationModule, bool) {
	m, ok := r.modules[moduleID]
	return m, ok
}

// ListModules returns info about all registered modules
func (r *ObligationsRegistry) ListModules() []RegulationInfo {
	var result []RegulationInfo
	for _, m := range r.modules {
		result = append(result, RegulationInfo{
			ID:          m.ID(),
			Name:        m.Name(),
			Description: m.Description(),
		})
	}
	// Sort by ID for consistent output
	sort.Slice(result, func(i, j int) bool {
		return result[i].ID < result[j].ID
	})
	return result
}

// EvaluateAll evaluates all registered modules against the given facts
func (r *ObligationsRegistry) EvaluateAll(tenantID uuid.UUID, facts *UnifiedFacts, orgName string) *ManagementObligationsOverview {
	overview := &ManagementObligationsOverview{
		ID:                    uuid.New(),
		TenantID:              tenantID,
		OrganizationName:      orgName,
		AssessmentDate:        time.Now(),
		CreatedAt:             time.Now(),
		UpdatedAt:             time.Now(),
		ApplicableRegulations: []ApplicableRegulation{},
		Obligations:           []Obligation{},
		RequiredControls:      []ObligationControl{},
		IncidentDeadlines:     []IncidentDeadline{},
	}

	// Track aggregated sanctions
	var maxFine string
	var personalLiability, criminalLiability bool
	var affectedRegulations []string

	// Evaluate each module
	for _, module := range r.modules {
		if module.IsApplicable(facts) {
			// Get classification
			classification := module.GetClassification(facts)

			// Derive obligations
			obligations := module.DeriveObligations(facts)

			// Derive controls
			controls := module.DeriveControls(facts)

			// Get incident deadlines
			incidentDeadlines := module.GetIncidentDeadlines(facts)

			// Add to applicable regulations
			overview.ApplicableRegulations = append(overview.ApplicableRegulations, ApplicableRegulation{
				ID:              module.ID(),
				Name:            module.Name(),
				Classification:  classification,
				Reason:          r.getApplicabilityReason(module, facts, classification),
				ObligationCount: len(obligations),
				ControlCount:    len(controls),
			})

			// Aggregate obligations
			overview.Obligations = append(overview.Obligations, obligations...)

			// Aggregate controls
			overview.RequiredControls = append(overview.RequiredControls, controls...)

			// Aggregate incident deadlines
			overview.IncidentDeadlines = append(overview.IncidentDeadlines, incidentDeadlines...)

			// Track sanctions
			for _, obl := range obligations {
				if obl.Sanctions != nil {
					if obl.Sanctions.MaxFine != "" && (maxFine == "" || len(obl.Sanctions.MaxFine) > len(maxFine)) {
						maxFine = obl.Sanctions.MaxFine
					}
					if obl.Sanctions.PersonalLiability {
						personalLiability = true
					}
					if obl.Sanctions.CriminalLiability {
						criminalLiability = true
					}
					if !containsString(affectedRegulations, module.ID()) {
						affectedRegulations = append(affectedRegulations, module.ID())
					}
				}
			}
		}
	}

	// Sort obligations by priority and deadline
	r.sortObligations(overview)

	// Build sanctions summary
	overview.SanctionsSummary = r.buildSanctionsSummary(maxFine, personalLiability, criminalLiability, affectedRegulations)

	// Generate executive summary
	overview.ExecutiveSummary = r.generateExecutiveSummary(overview)

	return overview
}

// EvaluateSingle evaluates a single module against the given facts
func (r *ObligationsRegistry) EvaluateSingle(moduleID string, facts *UnifiedFacts) (*ManagementObligationsOverview, error) {
	module, ok := r.modules[moduleID]
	if !ok {
		return nil, fmt.Errorf("module not found: %s", moduleID)
	}

	overview := &ManagementObligationsOverview{
		ID:                    uuid.New(),
		AssessmentDate:        time.Now(),
		CreatedAt:             time.Now(),
		UpdatedAt:             time.Now(),
		ApplicableRegulations: []ApplicableRegulation{},
		Obligations:           []Obligation{},
		RequiredControls:      []ObligationControl{},
		IncidentDeadlines:     []IncidentDeadline{},
	}

	if !module.IsApplicable(facts) {
		return overview, nil
	}

	classification := module.GetClassification(facts)
	obligations := module.DeriveObligations(facts)
	controls := module.DeriveControls(facts)
	incidentDeadlines := module.GetIncidentDeadlines(facts)

	overview.ApplicableRegulations = append(overview.ApplicableRegulations, ApplicableRegulation{
		ID:              module.ID(),
		Name:            module.Name(),
		Classification:  classification,
		Reason:          r.getApplicabilityReason(module, facts, classification),
		ObligationCount: len(obligations),
		ControlCount:    len(controls),
	})

	overview.Obligations = obligations
	overview.RequiredControls = controls
	overview.IncidentDeadlines = incidentDeadlines

	r.sortObligations(overview)
	overview.ExecutiveSummary = r.generateExecutiveSummary(overview)

	return overview, nil
}

// GetDecisionTree returns the decision tree for a specific module
func (r *ObligationsRegistry) GetDecisionTree(moduleID string) (*DecisionTree, error) {
	module, ok := r.modules[moduleID]
	if !ok {
		return nil, fmt.Errorf("module not found: %s", moduleID)
	}
	tree := module.GetDecisionTree()
	if tree == nil {
		return nil, fmt.Errorf("module %s does not have a decision tree", moduleID)
	}
	return tree, nil
}

// ============================================================================
// Helper Methods
// ============================================================================

func (r *ObligationsRegistry) getApplicabilityReason(module RegulationModule, facts *UnifiedFacts, classification string) string {
	switch module.ID() {
	case "nis2":
		if classification == string(NIS2EssentialEntity) {
			return "Besonders wichtige Einrichtung aufgrund von Sektor und Größe"
		} else if classification == string(NIS2ImportantEntity) {
			return "Wichtige Einrichtung aufgrund von Sektor und Größe"
		}
		return "NIS2-Richtlinie anwendbar"
	case "dsgvo":
		return "Verarbeitung personenbezogener Daten"
	case "ai_act":
		return "Einsatz von KI-Systemen"
	case "dora":
		return "Reguliertes Finanzunternehmen"
	default:
		return "Regulierung anwendbar"
	}
}

func (r *ObligationsRegistry) sortObligations(overview *ManagementObligationsOverview) {
	// Sort by priority (critical first), then by deadline
	priorityOrder := map[ObligationPriority]int{
		PriorityCritical: 0,
		PriorityHigh:     1,
		PriorityMedium:   2,
		PriorityLow:      3,
	}

	sort.Slice(overview.Obligations, func(i, j int) bool {
		// First by priority
		pi := priorityOrder[overview.Obligations[i].Priority]
		pj := priorityOrder[overview.Obligations[j].Priority]
		if pi != pj {
			return pi < pj
		}

		// Then by deadline (earlier first, nil last)
		di := overview.Obligations[i].Deadline
		dj := overview.Obligations[j].Deadline

		if di == nil && dj == nil {
			return false
		}
		if di == nil {
			return false
		}
		if dj == nil {
			return true
		}

		// For absolute deadlines, compare dates
		if di.Type == DeadlineAbsolute && dj.Type == DeadlineAbsolute {
			if di.Date != nil && dj.Date != nil {
				return di.Date.Before(*dj.Date)
			}
		}

		return false
	})
}

func (r *ObligationsRegistry) buildSanctionsSummary(maxFine string, personal, criminal bool, affected []string) SanctionsSummary {
	var summary string
	if personal && criminal {
		summary = "Hohe Bußgelder möglich. Persönliche Haftung der Geschäftsführung sowie strafrechtliche Konsequenzen bei Verstößen."
	} else if personal {
		summary = "Hohe Bußgelder möglich. Persönliche Haftung der Geschäftsführung bei Verstößen."
	} else if maxFine != "" {
		summary = fmt.Sprintf("Bußgelder bis zu %s bei Verstößen möglich.", maxFine)
	} else {
		summary = "Keine spezifischen Sanktionen identifiziert."
	}

	return SanctionsSummary{
		MaxFinancialRisk:      maxFine,
		PersonalLiabilityRisk: personal,
		CriminalLiabilityRisk: criminal,
		AffectedRegulations:   affected,
		Summary:               summary,
	}
}

func (r *ObligationsRegistry) generateExecutiveSummary(overview *ManagementObligationsOverview) ExecutiveSummary {
	summary := ExecutiveSummary{
		TotalRegulations:    len(overview.ApplicableRegulations),
		TotalObligations:    len(overview.Obligations),
		CriticalObligations: 0,
		UpcomingDeadlines:   0,
		OverdueObligations:  0,
		KeyRisks:            []string{},
		RecommendedActions:  []string{},
		ComplianceScore:     100, // Start at 100, deduct for gaps
	}

	now := time.Now()
	thirtyDaysFromNow := now.AddDate(0, 0, 30)

	for _, obl := range overview.Obligations {
		// Count critical
		if obl.Priority == PriorityCritical {
			summary.CriticalObligations++
			summary.ComplianceScore -= 10
		}

		// Count deadlines
		if obl.Deadline != nil && obl.Deadline.Type == DeadlineAbsolute && obl.Deadline.Date != nil {
			if obl.Deadline.Date.Before(now) {
				summary.OverdueObligations++
				summary.ComplianceScore -= 15
			} else if obl.Deadline.Date.Before(thirtyDaysFromNow) {
				summary.UpcomingDeadlines++
			}
		}
	}

	// Ensure score doesn't go below 0
	if summary.ComplianceScore < 0 {
		summary.ComplianceScore = 0
	}

	// Add key risks
	if summary.CriticalObligations > 0 {
		summary.KeyRisks = append(summary.KeyRisks, fmt.Sprintf("%d kritische Pflichten erfordern sofortige Aufmerksamkeit", summary.CriticalObligations))
	}
	if summary.OverdueObligations > 0 {
		summary.KeyRisks = append(summary.KeyRisks, fmt.Sprintf("%d Pflichten haben überfällige Fristen", summary.OverdueObligations))
	}
	if overview.SanctionsSummary.PersonalLiabilityRisk {
		summary.KeyRisks = append(summary.KeyRisks, "Persönliche Haftungsrisiken für die Geschäftsführung bestehen")
	}

	// Add recommended actions
	if summary.OverdueObligations > 0 {
		summary.RecommendedActions = append(summary.RecommendedActions, "Überfällige Pflichten priorisieren und umgehend adressieren")
	}
	if summary.CriticalObligations > 0 {
		summary.RecommendedActions = append(summary.RecommendedActions, "Kritische Pflichten in der nächsten Vorstandssitzung besprechen")
	}
	if len(overview.IncidentDeadlines) > 0 {
		summary.RecommendedActions = append(summary.RecommendedActions, "Incident-Response-Prozesse gemäß Meldefristen etablieren")
	}

	// Default action if no specific risks
	if len(summary.RecommendedActions) == 0 {
		summary.RecommendedActions = append(summary.RecommendedActions, "Regelmäßige Compliance-Reviews durchführen")
	}

	// Set next review date (3 months from now)
	nextReview := now.AddDate(0, 3, 0)
	summary.NextReviewDate = &nextReview

	return summary
}

// containsString checks if a slice contains a string
func containsString(slice []string, s string) bool {
	for _, item := range slice {
		if item == s {
			return true
		}
	}
	return false
}

// ============================================================================
// Grouping Methods
// ============================================================================

// GroupByRegulation groups obligations by their regulation ID
func (r *ObligationsRegistry) GroupByRegulation(obligations []Obligation) map[string][]Obligation {
	result := make(map[string][]Obligation)
	for _, obl := range obligations {
		result[obl.RegulationID] = append(result[obl.RegulationID], obl)
	}
	return result
}

// GroupByDeadline groups obligations by deadline timeframe
func (r *ObligationsRegistry) GroupByDeadline(obligations []Obligation) ObligationsByDeadlineResponse {
	result := ObligationsByDeadlineResponse{
		Overdue:     []Obligation{},
		ThisWeek:    []Obligation{},
		ThisMonth:   []Obligation{},
		NextQuarter: []Obligation{},
		Later:       []Obligation{},
		NoDeadline:  []Obligation{},
	}

	now := time.Now()
	oneWeek := now.AddDate(0, 0, 7)
	oneMonth := now.AddDate(0, 1, 0)
	threeMonths := now.AddDate(0, 3, 0)

	for _, obl := range obligations {
		if obl.Deadline == nil || obl.Deadline.Type != DeadlineAbsolute || obl.Deadline.Date == nil {
			result.NoDeadline = append(result.NoDeadline, obl)
			continue
		}

		deadline := *obl.Deadline.Date
		switch {
		case deadline.Before(now):
			result.Overdue = append(result.Overdue, obl)
		case deadline.Before(oneWeek):
			result.ThisWeek = append(result.ThisWeek, obl)
		case deadline.Before(oneMonth):
			result.ThisMonth = append(result.ThisMonth, obl)
		case deadline.Before(threeMonths):
			result.NextQuarter = append(result.NextQuarter, obl)
		default:
			result.Later = append(result.Later, obl)
		}
	}

	return result
}

// GroupByResponsible groups obligations by responsible role
func (r *ObligationsRegistry) GroupByResponsible(obligations []Obligation) map[ResponsibleRole][]Obligation {
	result := make(map[ResponsibleRole][]Obligation)
	for _, obl := range obligations {
		result[obl.Responsible] = append(result[obl.Responsible], obl)
	}
	return result
}