package scanner

import (
	"math"
	"regexp"
	"strings"

	"github.com/dd0c/drift-agent/pkg/models"
)

// Scrubber removes sensitive values from drift diffs before transmission.
type Scrubber struct {
	mode     string // "strict", "permissive", "off"
	patterns []*regexp.Regexp
}

func NewScrubber(mode string) *Scrubber {
	s := &Scrubber{mode: mode}
	if mode != "off" {
		s.patterns = []*regexp.Regexp{
			// AWS keys
			regexp.MustCompile(`(?i)(AKIA[0-9A-Z]{16})`),
			regexp.MustCompile(`(?i)(aws_secret_access_key|secret_key)\s*=\s*\S+`),
			// Generic secrets
			regexp.MustCompile(`(?i)(password|secret|token|api_key|private_key)\s*=\s*\S+`),
			// RSA/PEM keys
			regexp.MustCompile(`-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]*?-----END [A-Z ]*PRIVATE KEY-----`),
			// Base64 blobs > 40 chars (likely encoded secrets)
			regexp.MustCompile(`[A-Za-z0-9+/]{40,}={0,2}`),
		}
	}
	return s
}

// ScrubResource redacts sensitive attribute values from a drifted resource.
func (s *Scrubber) ScrubResource(r models.DriftedResource) models.DriftedResource {
	if s.mode == "off" {
		return r
	}

	scrubbed := r
	for i, diff := range scrubbed.Diffs {
		scrubbed.Diffs[i].OldValue = s.scrubValue(diff.AttributeName, diff.OldValue)
		scrubbed.Diffs[i].NewValue = s.scrubValue(diff.AttributeName, diff.NewValue)
	}
	return scrubbed
}

// ScrubString redacts sensitive patterns from an arbitrary string.
func (s *Scrubber) ScrubString(input string) string {
	if s.mode == "off" {
		return input
	}

	result := input
	for _, p := range s.patterns {
		result = p.ReplaceAllString(result, "[REDACTED]")
	}

	// Shannon entropy check for high-entropy strings (likely tokens/keys)
	if s.mode == "strict" {
		result = s.entropyRedact(result)
	}

	return result
}

func (s *Scrubber) scrubValue(attrName, value string) string {
	// Always redact known sensitive attribute names
	sensitiveAttrs := []string{
		"password", "secret", "token", "api_key", "private_key",
		"access_key", "secret_key", "connection_string", "certificate",
	}
	lower := strings.ToLower(attrName)
	for _, sa := range sensitiveAttrs {
		if strings.Contains(lower, sa) {
			return "[REDACTED]"
		}
	}

	return s.ScrubString(value)
}

// entropyRedact checks for high-entropy substrings (Shannon entropy > 3.5 bits/char).
func (s *Scrubber) entropyRedact(input string) string {
	words := strings.Fields(input)
	for i, word := range words {
		if len(word) > 20 && shannonEntropy(word) > 3.5 {
			words[i] = "[HIGH_ENTROPY_REDACTED]"
		}
	}
	return strings.Join(words, " ")
}

func shannonEntropy(s string) float64 {
	if len(s) == 0 {
		return 0
	}
	freq := make(map[rune]float64)
	for _, c := range s {
		freq[c]++
	}
	length := float64(len([]rune(s)))
	entropy := 0.0
	for _, count := range freq {
		p := count / length
		if p > 0 {
			entropy -= p * math.Log2(p)
		}
	}
	return entropy
}
