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

package iace

import (
	"context"
	"fmt"
	"strings"

	"github.com/breakpilot/ai-compliance-sdk/internal/llm"
	"github.com/breakpilot/ai-compliance-sdk/internal/ucca"
	"github.com/google/uuid"
)

// ============================================================================
// TechFileGenerator — LLM-based generation of technical file sections
// ============================================================================

// TechFileGenerator generates technical file section content using LLM and RAG.
type TechFileGenerator struct {
	llmRegistry *llm.ProviderRegistry
	ragClient   *ucca.LegalRAGClient
	store       *Store
}

// NewTechFileGenerator creates a new TechFileGenerator.
func NewTechFileGenerator(registry *llm.ProviderRegistry, ragClient *ucca.LegalRAGClient, store *Store) *TechFileGenerator {
	return &TechFileGenerator{
		llmRegistry: registry,
		ragClient:   ragClient,
		store:       store,
	}
}

// SectionGenerationContext holds all project data needed for LLM section generation.
type SectionGenerationContext struct {
	Project         *Project
	Components      []Component
	Hazards         []Hazard
	Assessments     map[uuid.UUID][]RiskAssessment // keyed by hazardID
	Mitigations     map[uuid.UUID][]Mitigation     // keyed by hazardID
	Classifications []RegulatoryClassification
	Evidence        []Evidence
	RAGContext      string // aggregated text from RAG search
}

// ============================================================================
// BuildSectionContext — loads all project data + RAG context
// ============================================================================

// BuildSectionContext loads project data and RAG context for a given section type.
func (g *TechFileGenerator) BuildSectionContext(ctx context.Context, projectID uuid.UUID, sectionType string) (*SectionGenerationContext, error) {
	// Load project
	project, err := g.store.GetProject(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load project: %w", err)
	}
	if project == nil {
		return nil, fmt.Errorf("project %s not found", projectID)
	}

	// Load components
	components, err := g.store.ListComponents(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load components: %w", err)
	}

	// Load hazards
	hazards, err := g.store.ListHazards(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load hazards: %w", err)
	}

	// Load assessments and mitigations per hazard
	assessments := make(map[uuid.UUID][]RiskAssessment)
	mitigations := make(map[uuid.UUID][]Mitigation)
	for _, h := range hazards {
		a, err := g.store.ListAssessments(ctx, h.ID)
		if err != nil {
			return nil, fmt.Errorf("load assessments for hazard %s: %w", h.ID, err)
		}
		assessments[h.ID] = a

		m, err := g.store.ListMitigations(ctx, h.ID)
		if err != nil {
			return nil, fmt.Errorf("load mitigations for hazard %s: %w", h.ID, err)
		}
		mitigations[h.ID] = m
	}

	// Load classifications
	classifications, err := g.store.GetClassifications(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load classifications: %w", err)
	}

	// Load evidence
	evidence, err := g.store.ListEvidence(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load evidence: %w", err)
	}

	// Perform RAG search for section-specific context
	ragContext := ""
	if g.ragClient != nil {
		ragQuery := buildRAGQuery(sectionType)
		results, ragErr := g.ragClient.SearchCollection(ctx, "bp_iace_libraries", ragQuery, nil, 5)
		if ragErr == nil && len(results) > 0 {
			var ragParts []string
			for _, r := range results {
				entry := fmt.Sprintf("[%s] %s", r.RegulationShort, truncateForPrompt(r.Text, 400))
				ragParts = append(ragParts, entry)
			}
			ragContext = strings.Join(ragParts, "\n\n")
		}
		// RAG failure is non-fatal — we proceed without context
	}

	return &SectionGenerationContext{
		Project:         project,
		Components:      components,
		Hazards:         hazards,
		Assessments:     assessments,
		Mitigations:     mitigations,
		Classifications: classifications,
		Evidence:        evidence,
		RAGContext:      ragContext,
	}, nil
}

// ============================================================================
// GenerateSection — main entry point
// ============================================================================

// GenerateSection generates the content for a technical file section using LLM.
// If LLM is unavailable, returns an enhanced placeholder with project data.
func (g *TechFileGenerator) GenerateSection(ctx context.Context, projectID uuid.UUID, sectionType string) (string, error) {
	sctx, err := g.BuildSectionContext(ctx, projectID, sectionType)
	if err != nil {
		return "", fmt.Errorf("build section context: %w", err)
	}

	// Build prompts
	systemPrompt := getSystemPrompt(sectionType)
	userPrompt := buildUserPrompt(sctx, sectionType)

	// Attempt LLM generation
	resp, err := g.llmRegistry.Chat(ctx, &llm.ChatRequest{
		Messages: []llm.Message{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userPrompt},
		},
		Temperature: 0.15,
		MaxTokens:   4096,
	})
	if err != nil {
		// LLM unavailable — return structured fallback with real project data
		return appendCrossRefIfApplicable(buildFallbackContent(sctx, sectionType), sctx, sectionType), nil
	}

	return appendCrossRefIfApplicable(resp.Message.Content, sctx, sectionType), nil
}

// appendCrossRefIfApplicable adds the international cross-reference appendix
// (DIN/ANSI/GB/JIS) to the "standards_applied" section. For other section
// types it returns content unchanged. The appendix is built deterministically
// from the in-process registry, so it is never hallucinated by the LLM.
func appendCrossRefIfApplicable(content string, sctx *SectionGenerationContext, sectionType string) string {
	if sectionType != SectionStandardsApplied {
		return content
	}
	normIDs := suggestNormIDsForProject(sctx)
	appendix := RenderCrossRefAppendix(normIDs)
	if appendix == "" {
		return content
	}
	return content + appendix
}

// suggestNormIDsForProject reuses the existing SuggestNorms heuristic to pick
// the norms most likely applicable to this project. We only need the IDs;
// the rest of the SuggestNorms output (scores, reasons) is discarded.
func suggestNormIDsForProject(sctx *SectionGenerationContext) []string {
	if sctx == nil || sctx.Project == nil {
		return nil
	}
	hazardCats := make([]string, 0, len(sctx.Hazards))
	seenCat := map[string]bool{}
	for _, h := range sctx.Hazards {
		if h.Category != "" && !seenCat[h.Category] {
			seenCat[h.Category] = true
			hazardCats = append(hazardCats, h.Category)
		}
	}
	result := SuggestNorms(sctx.Project.MachineType, hazardCats, nil)
	if result == nil {
		return nil
	}
	ids := make([]string, 0, result.Total)
	seenID := map[string]bool{}
	push := func(suggs []NormSuggestion) {
		for _, s := range suggs {
			if s.Norm.ID == "" || seenID[s.Norm.ID] {
				continue
			}
			seenID[s.Norm.ID] = true
			ids = append(ids, s.Norm.ID)
		}
	}
	push(result.ANorms)
	push(result.B1Norms)
	push(result.B2Norms)
	push(result.CNorms)
	return ids
}