package sanitize

import "regexp"

type rule struct {
	pattern     *regexp.Regexp
	replacement string
}

var rules = []rule{
	{regexp.MustCompile(`(?i)x-anthropic-billing-header:[^\n]*\n?`), ""},
	{regexp.MustCompile(`(?i)\bcc_version=[^\s;,\n]+[;,]?\s*`), ""},
	{regexp.MustCompile(`(?i)\bcc_entrypoint=[^\s;,\n]+[;,]?\s*`), ""},
	{regexp.MustCompile(`(?i)\bcch=[a-f0-9]+[;,]?\s*`), ""},
	{regexp.MustCompile(`\bClaude Code\b`), "Cursor"},
	{regexp.MustCompile(`(?i)Anthropic['']s official CLI for Claude`), "Cursor AI assistant"},
	{regexp.MustCompile(`\bAnthropic\b`), "Cursor"},
	{regexp.MustCompile(`(?i)anthropic\.com`), "cursor.com"},
	{regexp.MustCompile(`(?i)claude\.ai`), "cursor.sh"},
	{regexp.MustCompile(`^[;,\s]+`), ""},
}

func SanitizeText(text string) string {
	for _, r := range rules {
		text = r.pattern.ReplaceAllString(text, r.replacement)
	}
	return text
}

func SanitizeMessages(messages []interface{}) []interface{} {
	result := make([]interface{}, len(messages))
	for i, raw := range messages {
		if raw == nil {
			result[i] = raw
			continue
		}
		m, ok := raw.(map[string]interface{})
		if !ok {
			result[i] = raw
			continue
		}
		newMsg := make(map[string]interface{}, len(m))
		for k, v := range m {
			newMsg[k] = v
		}
		switch c := m["content"].(type) {
		case string:
			newMsg["content"] = SanitizeText(c)
		case []interface{}:
			newParts := make([]interface{}, len(c))
			for j, p := range c {
				if pm, ok := p.(map[string]interface{}); ok && pm["type"] == "text" {
					if t, ok := pm["text"].(string); ok {
						newPart := make(map[string]interface{}, len(pm))
						for k, v := range pm {
							newPart[k] = v
						}
						newPart["text"] = SanitizeText(t)
						newParts[j] = newPart
						continue
					}
				}
				newParts[j] = p
			}
			newMsg["content"] = newParts
		}
		result[i] = newMsg
	}
	return result
}

func SanitizeSystem(system interface{}) interface{} {
	switch v := system.(type) {
	case string:
		return SanitizeText(v)
	case []interface{}:
		result := make([]interface{}, len(v))
		for i, p := range v {
			if pm, ok := p.(map[string]interface{}); ok && pm["type"] == "text" {
				if t, ok := pm["text"].(string); ok {
					newPart := make(map[string]interface{}, len(pm))
					for k, val := range pm {
						newPart[k] = val
					}
					newPart["text"] = SanitizeText(t)
					result[i] = newPart
					continue
				}
			}
			result[i] = p
		}
		return result
	}
	return system
}
