breakpilot-compliance/ai-compliance-sdk/internal/training/block_generator.go

package training

import (
	"context"
	"fmt"
	"math"

	"github.com/google/uuid"
)

// BlockGenerator orchestrates the Controls → Training Modules pipeline
type BlockGenerator struct {
	store            *Store
	contentGenerator *ContentGenerator
}

// NewBlockGenerator creates a new block generator
func NewBlockGenerator(store *Store, contentGenerator *ContentGenerator) *BlockGenerator {
	return &BlockGenerator{
		store:            store,
		contentGenerator: contentGenerator,
	}
}

// Preview performs a dry run: loads matching controls, computes module split and roles
func (bg *BlockGenerator) Preview(ctx context.Context, configID uuid.UUID) (*PreviewBlockResponse, error) {
	config, err := bg.store.GetBlockConfig(ctx, configID)
	if err != nil {
		return nil, fmt.Errorf("load block config: %w", err)
	}
	if config == nil {
		return nil, fmt.Errorf("block config not found")
	}

	controls, err := bg.store.QueryCanonicalControls(ctx,
		config.DomainFilter, config.CategoryFilter,
		config.SeverityFilter, config.TargetAudienceFilter,
	)
	if err != nil {
		return nil, fmt.Errorf("query controls: %w", err)
	}

	maxPerModule := config.MaxControlsPerModule
	if maxPerModule <= 0 {
		maxPerModule = 20
	}
	moduleCount := int(math.Ceil(float64(len(controls)) / float64(maxPerModule)))
	if moduleCount == 0 && len(controls) > 0 {
		moduleCount = 1
	}

	roles := bg.deriveRoles(controls, config.TargetAudienceFilter)

	return &PreviewBlockResponse{
		ControlCount:  len(controls),
		ModuleCount:   moduleCount,
		Controls:      controls,
		ProposedRoles: roles,
	}, nil
}

// Generate executes the full pipeline: Controls → Modules → Links → CTM → Content
func (bg *BlockGenerator) Generate(ctx context.Context, configID uuid.UUID, req GenerateBlockRequest) (*GenerateBlockResponse, error) {
	config, err := bg.store.GetBlockConfig(ctx, configID)
	if err != nil {
		return nil, fmt.Errorf("load block config: %w", err)
	}
	if config == nil {
		return nil, fmt.Errorf("block config not found")
	}

	// 1. Load matching controls
	controls, err := bg.store.QueryCanonicalControls(ctx,
		config.DomainFilter, config.CategoryFilter,
		config.SeverityFilter, config.TargetAudienceFilter,
	)
	if err != nil {
		return nil, fmt.Errorf("query controls: %w", err)
	}

	if len(controls) == 0 {
		return &GenerateBlockResponse{}, nil
	}

	// 2. Chunk controls into module-sized groups
	maxPerModule := config.MaxControlsPerModule
	if maxPerModule <= 0 {
		maxPerModule = 20
	}
	chunks := chunkControls(controls, maxPerModule)

	// 3. Derive target roles for CTM
	roles := bg.deriveRoles(controls, config.TargetAudienceFilter)

	// 4. Count existing modules with this prefix for auto-numbering
	existingCount, err := bg.store.CountModulesWithPrefix(ctx, config.TenantID, config.ModuleCodePrefix)
	if err != nil {
		existingCount = 0
	}

	language := req.Language
	if language == "" {
		language = "de"
	}

	resp := &GenerateBlockResponse{}

	for i, chunk := range chunks {
		moduleNum := existingCount + i + 1
		moduleCode := fmt.Sprintf("%s-%02d", config.ModuleCodePrefix, moduleNum)

		// Build a descriptive title from the first few controls
		title := bg.buildModuleTitle(config, chunk, i+1, len(chunks))

		// a. Create TrainingModule
		module := &TrainingModule{
			TenantID:        config.TenantID,
			ModuleCode:      moduleCode,
			Title:           title,
			Description:     config.Description,
			RegulationArea:  config.RegulationArea,
			NIS2Relevant:    config.RegulationArea == RegulationNIS2,
			ISOControls:     bg.extractControlIDs(chunk),
			FrequencyType:   config.FrequencyType,
			ValidityDays:    365,
			RiskWeight:      2.0,
			ContentType:     "text",
			DurationMinutes: config.DurationMinutes,
			PassThreshold:   config.PassThreshold,
			IsActive:        true,
			SortOrder:       moduleNum,
		}

		if err := bg.store.CreateModule(ctx, module); err != nil {
			resp.Errors = append(resp.Errors, fmt.Sprintf("create module %s: %v", moduleCode, err))
			continue
		}
		resp.ModulesCreated++

		// b. Create control links (traceability)
		for j, ctrl := range chunk {
			link := &TrainingBlockControlLink{
				BlockConfigID:       config.ID,
				ModuleID:            module.ID,
				ControlID:           ctrl.ControlID,
				ControlTitle:        ctrl.Title,
				ControlObjective:    ctrl.Objective,
				ControlRequirements: ctrl.Requirements,
				SortOrder:           j,
			}
			if err := bg.store.CreateBlockControlLink(ctx, link); err != nil {
				resp.Errors = append(resp.Errors, fmt.Sprintf("link %s→%s: %v", moduleCode, ctrl.ControlID, err))
				continue
			}
			resp.ControlsLinked++
		}

		// c. Create CTM entries (target_audience → roles)
		if req.AutoMatrix {
			for _, role := range roles {
				entry := &TrainingMatrixEntry{
					TenantID:    config.TenantID,
					RoleCode:    role,
					ModuleID:    module.ID,
					IsMandatory: true,
					Priority:    1,
				}
				if err := bg.store.SetMatrixEntry(ctx, entry); err != nil {
					resp.Errors = append(resp.Errors, fmt.Sprintf("matrix %s→%s: %v", role, moduleCode, err))
					continue
				}
				resp.MatrixEntriesCreated++
			}
		}

		// d. Generate LLM content
		_, err := bg.contentGenerator.GenerateBlockContent(ctx, *module, chunk, language)
		if err != nil {
			resp.Errors = append(resp.Errors, fmt.Sprintf("content %s: %v", moduleCode, err))
			continue
		}
		resp.ContentGenerated++
	}

	// 5. Update last_generated_at
	bg.store.UpdateBlockConfigLastGenerated(ctx, config.ID)

	// 6. Audit log
	bg.store.LogAction(ctx, &AuditLogEntry{
		TenantID:   config.TenantID,
		Action:     AuditAction("block_generated"),
		EntityType: AuditEntityModule,
		Details: map[string]interface{}{
			"block_config_id":  config.ID.String(),
			"block_name":       config.Name,
			"modules_created":  resp.ModulesCreated,
			"controls_linked":  resp.ControlsLinked,
			"content_generated": resp.ContentGenerated,
		},
	})

	return resp, nil
}

// deriveRoles computes which CTM roles should receive the generated modules
func (bg *BlockGenerator) deriveRoles(controls []CanonicalControlSummary, audienceFilter string) []string {
	roleSet := map[string]bool{}

	// If a specific audience filter is set, use the mapping
	if audienceFilter != "" {
		if roles, ok := TargetAudienceRoleMapping[audienceFilter]; ok {
			for _, r := range roles {
				roleSet[r] = true
			}
		}
	}

	// Additionally derive roles from control categories
	for _, ctrl := range controls {
		if ctrl.Category != "" {
			if roles, ok := CategoryRoleMapping[ctrl.Category]; ok {
				for _, r := range roles {
					roleSet[r] = true
				}
			}
		}
		// Also check per-control target_audience
		if ctrl.TargetAudience != "" && audienceFilter == "" {
			if roles, ok := TargetAudienceRoleMapping[ctrl.TargetAudience]; ok {
				for _, r := range roles {
					roleSet[r] = true
				}
			}
		}
	}

	// If nothing derived, default to R9 (Alle Mitarbeiter)
	if len(roleSet) == 0 {
		roleSet[RoleR9] = true
	}

	roles := make([]string, 0, len(roleSet))
	for r := range roleSet {
		roles = append(roles, r)
	}
	return roles
}

// buildModuleTitle creates a descriptive module title
func (bg *BlockGenerator) buildModuleTitle(config *TrainingBlockConfig, controls []CanonicalControlSummary, partNum, totalParts int) string {
	base := config.Name
	if totalParts > 1 {
		base = fmt.Sprintf("%s (Teil %d/%d)", config.Name, partNum, totalParts)
	}
	return base
}

// extractControlIDs returns the control IDs from a slice of controls
func (bg *BlockGenerator) extractControlIDs(controls []CanonicalControlSummary) []string {
	ids := make([]string, len(controls))
	for i, c := range controls {
		ids[i] = c.ControlID
	}
	return ids
}

// chunkControls splits controls into groups of maxSize
func chunkControls(controls []CanonicalControlSummary, maxSize int) [][]CanonicalControlSummary {
	if maxSize <= 0 {
		maxSize = 20
	}

	var chunks [][]CanonicalControlSummary
	for i := 0; i < len(controls); i += maxSize {
		end := i + maxSize
		if end > len(controls) {
			end = len(controls)
		}
		chunks = append(chunks, controls[i:end])
	}
	return chunks
}