package ucca

import (
	"context"
	"fmt"
	"strings"
	"time"
)

// GetLegalContextForAssessment retrieves relevant legal context for an assessment.
func (c *LegalRAGClient) GetLegalContextForAssessment(ctx context.Context, assessment *Assessment) (*LegalContext, error) {
	queryParts := []string{}

	if assessment.Domain != "" {
		queryParts = append(queryParts, fmt.Sprintf("KI-Anwendung im Bereich %s", assessment.Domain))
	}

	if assessment.Intake.DataTypes.Article9Data {
		queryParts = append(queryParts, "besondere Kategorien personenbezogener Daten Art. 9 DSGVO")
	}
	if assessment.Intake.DataTypes.PersonalData {
		queryParts = append(queryParts, "personenbezogene Daten")
	}
	if assessment.Intake.DataTypes.MinorData {
		queryParts = append(queryParts, "Daten von Minderjährigen")
	}

	if assessment.Intake.Purpose.EvaluationScoring {
		queryParts = append(queryParts, "automatisierte Bewertung Scoring")
	}
	if assessment.Intake.Purpose.DecisionMaking {
		queryParts = append(queryParts, "automatisierte Entscheidung Art. 22 DSGVO")
	}
	if assessment.Intake.Purpose.Profiling {
		queryParts = append(queryParts, "Profiling")
	}

	if assessment.DSFARecommended {
		queryParts = append(queryParts, "Datenschutz-Folgenabschätzung Art. 35 DSGVO")
	}
	if assessment.Art22Risk {
		queryParts = append(queryParts, "automatisierte Einzelentscheidung rechtliche Wirkung")
	}

	query := strings.Join(queryParts, " ")
	if query == "" {
		query = "DSGVO Anforderungen KI-System Datenschutz"
	}

	regulationIDs := c.determineRelevantRegulations(assessment)

	results, err := c.Search(ctx, query, regulationIDs, 5)
	if err != nil {
		return nil, err
	}

	regSet := make(map[string]bool)
	for _, r := range results {
		regSet[r.RegulationCode] = true
	}

	regulations := make([]string, 0, len(regSet))
	for r := range regSet {
		regulations = append(regulations, r)
	}

	articles := make([]string, 0)
	for _, r := range results {
		if len(r.Pages) > 0 {
			key := fmt.Sprintf("%s S. %v", r.RegulationShort, r.Pages)
			articles = append(articles, key)
		}
	}

	return &LegalContext{
		Query:            query,
		Results:          results,
		RelevantArticles: articles,
		Regulations:      regulations,
		GeneratedAt:      time.Now().UTC(),
	}, nil
}

// determineRelevantRegulations determines which regulations to search based on the assessment.
func (c *LegalRAGClient) determineRelevantRegulations(assessment *Assessment) []string {
	ids := []string{"eu_2016_679"}

	for _, rule := range assessment.TriggeredRules {
		gdprRef := rule.GDPRRef
		if strings.Contains(gdprRef, "AI Act") || strings.Contains(gdprRef, "KI-VO") {
			if !contains(ids, "eu_2024_1689") {
				ids = append(ids, "eu_2024_1689")
			}
		}
		if strings.Contains(gdprRef, "NIS2") || strings.Contains(gdprRef, "NIS-2") {
			if !contains(ids, "eu_2022_2555") {
				ids = append(ids, "eu_2022_2555")
			}
		}
		if strings.Contains(gdprRef, "CRA") || strings.Contains(gdprRef, "Cyber Resilience") {
			if !contains(ids, "eu_2024_2847") {
				ids = append(ids, "eu_2024_2847")
			}
		}
		if strings.Contains(gdprRef, "Maschinenverordnung") || strings.Contains(gdprRef, "Machinery") {
			if !contains(ids, "eu_2023_1230") {
				ids = append(ids, "eu_2023_1230")
			}
		}
	}

	for _, ctrl := range assessment.RequiredControls {
		if strings.HasPrefix(ctrl.ID, "AI-") {
			if !contains(ids, "eu_2024_1689") {
				ids = append(ids, "eu_2024_1689")
			}
			break
		}
	}

	for _, ctrl := range assessment.RequiredControls {
		if strings.HasPrefix(ctrl.ID, "CRYPTO-") || strings.HasPrefix(ctrl.ID, "IAM-") || strings.HasPrefix(ctrl.ID, "SEC-") {
			if !contains(ids, "eu_2022_2555") {
				ids = append(ids, "eu_2022_2555")
			}
			if !contains(ids, "eu_2024_2847") {
				ids = append(ids, "eu_2024_2847")
			}
			break
		}
	}

	return ids
}