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

package handlers

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

	"github.com/breakpilot/ai-compliance-sdk/internal/llm"
	"github.com/breakpilot/ai-compliance-sdk/internal/ucca"
	"github.com/gin-gonic/gin"
)

// WizardAskRequest represents a question to the Legal Assistant
type WizardAskRequest struct {
	Question    string                 `json:"question" binding:"required"`
	StepNumber  int                    `json:"step_number"`
	FieldID     string                 `json:"field_id,omitempty"`
	CurrentData map[string]interface{} `json:"current_data,omitempty"`
}

// WizardAskResponse represents the Legal Assistant response
type WizardAskResponse struct {
	Answer        string        `json:"answer"`
	Sources       []LegalSource `json:"sources,omitempty"`
	RelatedFields []string      `json:"related_fields,omitempty"`
	GeneratedAt   time.Time     `json:"generated_at"`
	Model         string        `json:"model"`
}

// LegalSource represents a legal reference used in the answer
type LegalSource struct {
	Regulation string `json:"regulation"`
	Article    string `json:"article,omitempty"`
	Text       string `json:"text,omitempty"`
}

// AskWizardQuestion handles legal questions from the wizard
func (h *UCCAHandlers) AskWizardQuestion(c *gin.Context) {
	var req WizardAskRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	ragQuery := buildWizardRAGQuery(req)

	var legalResults []ucca.LegalSearchResult
	var sources []LegalSource
	if h.legalRAGClient != nil {
		results, err := h.legalRAGClient.Search(c.Request.Context(), ragQuery, nil, 5)
		if err != nil {
			fmt.Printf("Warning: Legal RAG search failed: %v\n", err)
		} else {
			legalResults = results
			sources = make([]LegalSource, len(results))
			for i, r := range results {
				sources[i] = LegalSource{
					Regulation: r.RegulationName,
					Article:    r.Article,
					Text:       truncateText(r.Text, 200),
				}
			}
		}
	}

	prompt := buildWizardAssistantPrompt(req, legalResults)
	systemPrompt := buildWizardSystemPrompt(req.StepNumber)

	chatReq := &llm.ChatRequest{
		Messages: []llm.Message{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: prompt},
		},
		MaxTokens:   1024,
		Temperature: 0.3,
	}
	response, err := h.providerRegistry.Chat(c.Request.Context(), chatReq)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "LLM call failed: " + err.Error()})
		return
	}

	relatedFields := identifyRelatedFields(req.Question)

	c.JSON(http.StatusOK, WizardAskResponse{
		Answer:        response.Message.Content,
		Sources:       sources,
		RelatedFields: relatedFields,
		GeneratedAt:   time.Now().UTC(),
		Model:         response.Model,
	})
}

// GetWizardSchema returns the wizard schema for the frontend
func (h *UCCAHandlers) GetWizardSchema(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{
		"version":      "1.1",
		"total_steps":  8,
		"default_mode": "simple",
		"legal_assistant": gin.H{
			"enabled":  true,
			"endpoint": "/sdk/v1/ucca/wizard/ask",
			"max_tokens": 1024,
			"example_questions": []string{
				"Was sind personenbezogene Daten?",
				"Was ist der Unterschied zwischen AVV und SCC?",
				"Brauche ich ein TIA?",
				"Was bedeutet Profiling?",
				"Was ist Art. 9 DSGVO?",
				"Wann brauche ich eine DSFA?",
				"Was ist das Data Privacy Framework?",
			},
		},
		"steps": []gin.H{
			{"number": 1, "title": "Grundlegende Informationen", "icon": "info"},
			{"number": 2, "title": "Welche Daten werden verarbeitet?", "icon": "database"},
			{"number": 3, "title": "Wofür wird die KI eingesetzt?", "icon": "target"},
			{"number": 4, "title": "Wo läuft die KI?", "icon": "server"},
			{"number": 5, "title": "Internationaler Datentransfer", "icon": "globe"},
			{"number": 6, "title": "KI-Modell und Training", "icon": "brain"},
			{"number": 7, "title": "Verträge & Compliance", "icon": "file-contract"},
			{"number": 8, "title": "Automatisierung & Kontrolle", "icon": "user-check"},
		},
	})
}

// buildWizardRAGQuery creates an optimized query for Legal RAG search
func buildWizardRAGQuery(req WizardAskRequest) string {
	query := req.Question

	stepContext := map[int]string{
		1: "KI-Anwendung Use Case",
		2: "personenbezogene Daten Datenkategorien DSGVO Art. 4 Art. 9",
		3: "Verarbeitungszweck Profiling Scoring automatisierte Entscheidung Art. 22",
		4: "Hosting Cloud On-Premises Auftragsverarbeitung",
		5: "Standardvertragsklauseln SCC Drittlandtransfer TIA Transfer Impact Assessment Art. 44 Art. 46",
		6: "KI-Modell Training RAG Finetuning",
		7: "Auftragsverarbeitungsvertrag AVV DSFA Verarbeitungsverzeichnis Art. 28 Art. 30 Art. 35",
		8: "Automatisierung Human-in-the-Loop Art. 22 AI Act",
	}

	if context, ok := stepContext[req.StepNumber]; ok {
		query = query + " " + context
	}

	return query
}

// buildWizardSystemPrompt creates the system prompt for the Legal Assistant
func buildWizardSystemPrompt(stepNumber int) string {
	basePrompt := `Du bist ein freundlicher Rechtsassistent, der Nutzern hilft,
datenschutzrechtliche Begriffe und Anforderungen zu verstehen.

DEINE AUFGABE:
- Erkläre rechtliche Begriffe in einfacher, verständlicher Sprache
- Beantworte Fragen zum aktuellen Wizard-Schritt
- Hilf dem Nutzer, die richtigen Antworten im Wizard zu geben
- Verweise auf relevante Rechtsquellen (DSGVO-Artikel, etc.)

WICHTIGE REGELN:
- Antworte IMMER auf Deutsch
- Verwende einfache Sprache, keine Juristensprache
- Gib konkrete Beispiele wenn möglich
- Bei Unsicherheit empfehle die Rücksprache mit einem Datenschutzbeauftragten
- Du darfst KEINE Rechtsberatung geben, nur erklären

ANTWORT-FORMAT:
- Kurz und prägnant (max. 3-4 Sätze für einfache Fragen)
- Strukturiert mit Aufzählungen bei komplexen Themen
- Immer mit Quellenangabe am Ende (z.B. "Siehe: DSGVO Art. 9")`

	stepContexts := map[int]string{
		1: "\n\nKONTEXT: Der Nutzer befindet sich im ersten Schritt und gibt grundlegende Informationen zum KI-Vorhaben ein.",
		2: "\n\nKONTEXT: Der Nutzer gibt an, welche Datenarten verarbeitet werden. Erkläre die Unterschiede zwischen personenbezogenen Daten, Art. 9 Daten (besondere Kategorien), biometrischen Daten, etc.",
		3: "\n\nKONTEXT: Der Nutzer gibt den Verarbeitungszweck an. Erkläre Begriffe wie Profiling, Scoring, systematische Überwachung, automatisierte Entscheidungen mit rechtlicher Wirkung.",
		4: "\n\nKONTEXT: Der Nutzer gibt Hosting-Informationen an. Erkläre Cloud vs. On-Premises, wann Drittlandtransfer vorliegt, Unterschiede zwischen EU/EWR und Drittländern.",
		5: "\n\nKONTEXT: Der Nutzer beantwortet Fragen zu SCC und TIA. Erkläre Standardvertragsklauseln (SCC), Transfer Impact Assessment (TIA), das Data Privacy Framework (DPF), und wann welche Instrumente erforderlich sind.",
		6: "\n\nKONTEXT: Der Nutzer gibt KI-Modell-Informationen an. Erkläre RAG vs. Training/Finetuning, warum Training mit personenbezogenen Daten problematisch ist, und welche Opt-Out-Klauseln wichtig sind.",
		7: "\n\nKONTEXT: Der Nutzer beantwortet Fragen zu Verträgen. Erkläre den Auftragsverarbeitungsvertrag (AVV), die Datenschutz-Folgenabschätzung (DSFA), das Verarbeitungsverzeichnis (VVT), und wann diese erforderlich sind.",
		8: "\n\nKONTEXT: Der Nutzer gibt den Automatisierungsgrad an. Erkläre Human-in-the-Loop, Art. 22 DSGVO (automatisierte Einzelentscheidungen), und die Anforderungen des AI Acts.",
	}

	if context, ok := stepContexts[stepNumber]; ok {
		basePrompt += context
	}

	return basePrompt
}

// buildWizardAssistantPrompt creates the user prompt with legal context
func buildWizardAssistantPrompt(req WizardAskRequest, legalResults []ucca.LegalSearchResult) string {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("FRAGE DES NUTZERS:\n%s\n\n", req.Question))

	if len(legalResults) > 0 {
		buf.WriteString("RELEVANTE RECHTSGRUNDLAGEN (aus unserer Bibliothek):\n\n")
		for i, result := range legalResults {
			buf.WriteString(fmt.Sprintf("%d. %s", i+1, result.RegulationName))
			if result.Article != "" {
				buf.WriteString(fmt.Sprintf(" - Art. %s", result.Article))
				if result.Paragraph != "" {
					buf.WriteString(fmt.Sprintf(" Abs. %s", result.Paragraph))
				}
			}
			buf.WriteString("\n")
			buf.WriteString(fmt.Sprintf("   %s\n\n", truncateText(result.Text, 300)))
		}
	}

	if req.FieldID != "" {
		buf.WriteString(fmt.Sprintf("AKTUELLES FELD: %s\n\n", req.FieldID))
	}

	buf.WriteString("Bitte beantworte die Frage kurz und verständlich. Verwende die angegebenen Rechtsgrundlagen als Referenz.")

	return buf.String()
}

// identifyRelatedFields identifies wizard fields related to the question
func identifyRelatedFields(question string) []string {
	question = strings.ToLower(question)
	var related []string

	keywordMapping := map[string][]string{
		"personenbezogen":          {"data_types.personal_data"},
		"art. 9":                   {"data_types.article_9_data"},
		"sensibel":                 {"data_types.article_9_data"},
		"gesundheit":               {"data_types.article_9_data"},
		"minderjährig":             {"data_types.minor_data"},
		"kinder":                   {"data_types.minor_data"},
		"biometrisch":              {"data_types.biometric_data"},
		"gesicht":                  {"data_types.biometric_data"},
		"kennzeichen":              {"data_types.license_plates"},
		"standort":                 {"data_types.location_data"},
		"gps":                      {"data_types.location_data"},
		"profiling":                {"purpose.profiling"},
		"scoring":                  {"purpose.evaluation_scoring"},
		"überwachung":              {"processing.systematic_monitoring"},
		"automatisch":              {"outputs.decision_with_legal_effect", "automation"},
		"entscheidung":             {"outputs.decision_with_legal_effect"},
		"cloud":                    {"hosting.type", "hosting.region"},
		"on-premises":              {"hosting.type"},
		"lokal":                    {"hosting.type"},
		"scc":                      {"contracts.scc.present", "contracts.scc.version"},
		"standardvertrags":         {"contracts.scc.present"},
		"drittland":                {"hosting.region", "provider.location"},
		"usa":                      {"hosting.region", "provider.location", "provider.dpf_certified"},
		"transfer":                 {"hosting.region", "contracts.tia.present"},
		"tia":                      {"contracts.tia.present", "contracts.tia.result"},
		"dpf":                      {"provider.dpf_certified"},
		"data privacy":             {"provider.dpf_certified"},
		"avv":                      {"contracts.avv.present"},
		"auftragsverarbeitung":     {"contracts.avv.present"},
		"dsfa":                     {"governance.dsfa_completed"},
		"folgenabschätzung":        {"governance.dsfa_completed"},
		"verarbeitungsverzeichnis": {"governance.vvt_entry"},
		"training":                 {"model_usage.training", "provider.uses_data_for_training"},
		"finetuning":               {"model_usage.training"},
		"rag":                      {"model_usage.rag"},
		"human":                    {"processing.human_oversight"},
		"aufsicht":                 {"processing.human_oversight"},
	}

	seen := make(map[string]bool)
	for keyword, fields := range keywordMapping {
		if strings.Contains(question, keyword) {
			for _, field := range fields {
				if !seen[field] {
					related = append(related, field)
					seen[field] = true
				}
			}
		}
	}

	return related
}