breakpilot-compliance/ai-compliance-sdk/internal/api/handlers/drafting_handlers.go

package handlers

import (
	"fmt"
	"net/http"
	"time"

	"github.com/breakpilot/ai-compliance-sdk/internal/audit"
	"github.com/breakpilot/ai-compliance-sdk/internal/llm"
	"github.com/breakpilot/ai-compliance-sdk/internal/rbac"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
)

// DraftingHandlers handles Drafting Engine API endpoints
type DraftingHandlers struct {
	accessGate   *llm.AccessGate
	registry     *llm.ProviderRegistry
	piiDetector  *llm.PIIDetector
	auditStore   *audit.Store
	trailBuilder *audit.TrailBuilder
}

// NewDraftingHandlers creates new Drafting Engine handlers
func NewDraftingHandlers(
	accessGate *llm.AccessGate,
	registry *llm.ProviderRegistry,
	piiDetector *llm.PIIDetector,
	auditStore *audit.Store,
	trailBuilder *audit.TrailBuilder,
) *DraftingHandlers {
	return &DraftingHandlers{
		accessGate:   accessGate,
		registry:     registry,
		piiDetector:  piiDetector,
		auditStore:   auditStore,
		trailBuilder: trailBuilder,
	}
}

// ---------------------------------------------------------------------------
// Request/Response Types
// ---------------------------------------------------------------------------

// DraftDocumentRequest represents a request to generate a compliance document draft
type DraftDocumentRequest struct {
	DocumentType string                 `json:"document_type" binding:"required"`
	ScopeLevel   string                 `json:"scope_level" binding:"required"`
	Context      map[string]interface{} `json:"context"`
	Instructions string                 `json:"instructions"`
	Model        string                 `json:"model"`
}

// ValidateDocumentRequest represents a request to validate document consistency
type ValidateDocumentRequest struct {
	DocumentType      string                 `json:"document_type" binding:"required"`
	DraftContent      string                 `json:"draft_content"`
	ValidationContext map[string]interface{} `json:"validation_context"`
	Model             string                 `json:"model"`
}

// DraftHistoryEntry represents a single audit trail entry for drafts
type DraftHistoryEntry struct {
	ID                   string    `json:"id"`
	UserID               string    `json:"user_id"`
	TenantID             string    `json:"tenant_id"`
	DocumentType         string    `json:"document_type"`
	ScopeLevel           string    `json:"scope_level"`
	Operation            string    `json:"operation"`
	ConstraintsRespected bool      `json:"constraints_respected"`
	TokensUsed           int       `json:"tokens_used"`
	CreatedAt            time.Time `json:"created_at"`
}

// ---------------------------------------------------------------------------
// Handlers
// ---------------------------------------------------------------------------

// DraftDocument handles document draft generation via LLM with constraint validation
func (h *DraftingHandlers) DraftDocument(c *gin.Context) {
	var req DraftDocumentRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	userID := rbac.GetUserID(c)
	tenantID := rbac.GetTenantID(c)
	namespaceID := rbac.GetNamespaceID(c)

	if userID == uuid.Nil || tenantID == uuid.Nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication required"})
		return
	}

	// Validate scope level
	validLevels := map[string]bool{"L1": true, "L2": true, "L3": true, "L4": true}
	if !validLevels[req.ScopeLevel] {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid scope_level, must be L1-L4"})
		return
	}

	// Validate document type
	validTypes := map[string]bool{
		"vvt": true, "tom": true, "dsfa": true, "dsi": true, "lf": true,
		"av_vertrag": true, "betroffenenrechte": true, "einwilligung": true,
		"daten_transfer": true, "datenpannen": true, "vertragsmanagement": true,
		"schulung": true, "audit_log": true, "risikoanalyse": true,
		"notfallplan": true, "zertifizierung": true, "datenschutzmanagement": true,
		"iace_ce_assessment": true,
	}
	if !validTypes[req.DocumentType] {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid document_type"})
		return
	}

	// Build system prompt for drafting
	systemPrompt := fmt.Sprintf(
		`Du bist ein DSGVO-Compliance-Experte. Erstelle einen strukturierten Entwurf fuer Dokument "%s" auf Level %s.
Antworte NUR im JSON-Format mit einem "sections" Array.
Jede Section hat: id, title, content, schemaField.
Halte die Tiefe strikt am vorgegebenen Level.
Markiere fehlende Informationen mit [PLATZHALTER: Beschreibung].
Sprache: Deutsch.`,
		req.DocumentType, req.ScopeLevel,
	)

	userPrompt := "Erstelle den Dokumententwurf."
	if req.Instructions != "" {
		userPrompt = req.Instructions
	}

	// Detect PII in context
	contextStr := fmt.Sprintf("%v", req.Context)
	dataCategories := h.piiDetector.DetectDataCategories(contextStr)

	// Process through access gate
	chatReq := &llm.ChatRequest{
		Model: req.Model,
		Messages: []llm.Message{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userPrompt},
		},
		MaxTokens:   16384,
		Temperature: 0.15,
	}

	gatedReq, err := h.accessGate.ProcessChatRequest(
		c.Request.Context(),
		userID, tenantID, namespaceID,
		chatReq, dataCategories,
	)
	if err != nil {
		h.logDraftAudit(c, userID, tenantID, req.DocumentType, req.ScopeLevel, "draft", false, 0, err.Error())
		c.JSON(http.StatusForbidden, gin.H{
			"error":   "access_denied",
			"message": err.Error(),
		})
		return
	}

	// Execute the request
	resp, err := h.accessGate.ExecuteChat(c.Request.Context(), gatedReq)
	if err != nil {
		h.logDraftAudit(c, userID, tenantID, req.DocumentType, req.ScopeLevel, "draft", false, 0, err.Error())
		c.JSON(http.StatusInternalServerError, gin.H{
			"error":   "llm_error",
			"message": err.Error(),
		})
		return
	}

	tokensUsed := 0
	if resp.Usage.TotalTokens > 0 {
		tokensUsed = resp.Usage.TotalTokens
	}

	// Log successful draft
	h.logDraftAudit(c, userID, tenantID, req.DocumentType, req.ScopeLevel, "draft", true, tokensUsed, "")

	c.JSON(http.StatusOK, gin.H{
		"document_type": req.DocumentType,
		"scope_level":   req.ScopeLevel,
		"content":       resp.Message.Content,
		"model":         resp.Model,
		"provider":      resp.Provider,
		"tokens_used":   tokensUsed,
		"pii_detected":  gatedReq.PIIDetected,
	})
}

// ValidateDocument handles document cross-consistency validation
func (h *DraftingHandlers) ValidateDocument(c *gin.Context) {
	var req ValidateDocumentRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	userID := rbac.GetUserID(c)
	tenantID := rbac.GetTenantID(c)
	namespaceID := rbac.GetNamespaceID(c)

	if userID == uuid.Nil || tenantID == uuid.Nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication required"})
		return
	}

	// Build validation prompt
	systemPrompt := `Du bist ein DSGVO-Compliance-Validator. Pruefe die Konsistenz und Vollstaendigkeit.
Antworte NUR im JSON-Format:
{
  "passed": boolean,
  "errors": [{"id": string, "severity": "error", "title": string, "description": string, "documentType": string, "legalReference": string}],
  "warnings": [{"id": string, "severity": "warning", "title": string, "description": string}],
  "suggestions": [{"id": string, "severity": "suggestion", "title": string, "description": string, "suggestion": string}]
}`

	validationPrompt := fmt.Sprintf("Validiere Dokument '%s'.\nInhalt:\n%s\nKontext:\n%v",
		req.DocumentType, req.DraftContent, req.ValidationContext)

	// Detect PII
	dataCategories := h.piiDetector.DetectDataCategories(req.DraftContent)

	chatReq := &llm.ChatRequest{
		Model: req.Model,
		Messages: []llm.Message{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: validationPrompt},
		},
		MaxTokens:   8192,
		Temperature: 0.1,
	}

	gatedReq, err := h.accessGate.ProcessChatRequest(
		c.Request.Context(),
		userID, tenantID, namespaceID,
		chatReq, dataCategories,
	)
	if err != nil {
		h.logDraftAudit(c, userID, tenantID, req.DocumentType, "", "validate", false, 0, err.Error())
		c.JSON(http.StatusForbidden, gin.H{
			"error":   "access_denied",
			"message": err.Error(),
		})
		return
	}

	resp, err := h.accessGate.ExecuteChat(c.Request.Context(), gatedReq)
	if err != nil {
		h.logDraftAudit(c, userID, tenantID, req.DocumentType, "", "validate", false, 0, err.Error())
		c.JSON(http.StatusInternalServerError, gin.H{
			"error":   "llm_error",
			"message": err.Error(),
		})
		return
	}

	tokensUsed := 0
	if resp.Usage.TotalTokens > 0 {
		tokensUsed = resp.Usage.TotalTokens
	}

	h.logDraftAudit(c, userID, tenantID, req.DocumentType, "", "validate", true, tokensUsed, "")

	c.JSON(http.StatusOK, gin.H{
		"document_type": req.DocumentType,
		"validation":    resp.Message.Content,
		"model":         resp.Model,
		"provider":      resp.Provider,
		"tokens_used":   tokensUsed,
		"pii_detected":  gatedReq.PIIDetected,
	})
}

// GetDraftHistory returns the audit trail of all drafting operations for a tenant
func (h *DraftingHandlers) GetDraftHistory(c *gin.Context) {
	userID := rbac.GetUserID(c)
	tenantID := rbac.GetTenantID(c)

	if userID == uuid.Nil || tenantID == uuid.Nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication required"})
		return
	}

	// Query audit store for drafting operations
	entries, _, err := h.auditStore.QueryGeneralAuditEntries(c.Request.Context(), &audit.GeneralAuditFilter{
		TenantID:     tenantID,
		ResourceType: "compliance_document",
		Limit:        50,
	})
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to query draft history"})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"history": entries,
		"total":   len(entries),
	})
}

// ---------------------------------------------------------------------------
// Audit Logging
// ---------------------------------------------------------------------------

func (h *DraftingHandlers) logDraftAudit(
	c *gin.Context,
	userID, tenantID uuid.UUID,
	documentType, scopeLevel, operation string,
	constraintsRespected bool,
	tokensUsed int,
	errorMsg string,
) {
	newValues := map[string]any{
		"document_type":        documentType,
		"scope_level":          scopeLevel,
		"constraints_respected": constraintsRespected,
		"tokens_used":          tokensUsed,
	}
	if errorMsg != "" {
		newValues["error"] = errorMsg
	}

	entry := h.trailBuilder.NewGeneralEntry().
		WithTenant(tenantID).
		WithUser(userID).
		WithAction("drafting_engine." + operation).
		WithResource("compliance_document", nil).
		WithNewValues(newValues).
		WithClient(c.ClientIP(), c.GetHeader("User-Agent"))

	go func() {
		entry.Save(c.Request.Context())
	}()
}