package proxy

import (
	"strings"
	"testing"

	"github.com/fujin/anthropic-proxy/internal/config"
)

func TestNewSanitizer_Empty(t *testing.T) {
	s := NewSanitizer(config.SanitizeConfig{})
	if len(s.toolsForward) != 0 {
		t.Errorf("expected empty toolsForward, got %d entries", len(s.toolsForward))
	}
	if len(s.toolsReverse) != 0 {
		t.Errorf("expected empty toolsReverse, got %d entries", len(s.toolsReverse))
	}
	if s.systemRules != nil {
		t.Errorf("expected nil systemRules")
	}
	if s.bodyRules != nil {
		t.Errorf("expected nil bodyRules")
	}
}

func TestNewSanitizer_WithTools(t *testing.T) {
	cfg := config.SanitizeConfig{
		Tools: []config.RenameRule{
			{From: "old_tool", To: "new_tool"},
			{From: "another", To: "replaced"},
		},
	}
	s := NewSanitizer(cfg)
	if got := s.toolsForward["old_tool"]; got != "new_tool" {
		t.Errorf("toolsForward[old_tool] = %q, want %q", got, "new_tool")
	}
	if got := s.toolsReverse["new_tool"]; got != "old_tool" {
		t.Errorf("toolsReverse[new_tool] = %q, want %q", got, "old_tool")
	}
	if got := s.toolsForward["another"]; got != "replaced" {
		t.Errorf("toolsForward[another] = %q, want %q", got, "replaced")
	}
	if got := s.toolsReverse["replaced"]; got != "another" {
		t.Errorf("toolsReverse[replaced] = %q, want %q", got, "another")
	}
}

func TestNewSanitizer_WithSystemAndBodyRules(t *testing.T) {
	cfg := config.SanitizeConfig{
		System: []config.ReplaceRule{{Match: "foo", Replace: "bar"}},
		Body:   []config.ReplaceRule{{Match: "baz", Replace: "qux"}},
	}
	s := NewSanitizer(cfg)
	if len(s.systemRules) != 1 || s.systemRules[0].Match != "foo" {
		t.Errorf("systemRules not set correctly")
	}
	if len(s.bodyRules) != 1 || s.bodyRules[0].Match != "baz" {
		t.Errorf("bodyRules not set correctly")
	}
}

func TestRenameTools(t *testing.T) {
	tests := []struct {
		name     string
		forward  map[string]string
		body     string
		expected string
	}{
		{
			name:     "empty map returns body unchanged",
			forward:  map[string]string{},
			body:     `{"tools":[{"name":"my_tool"}]}`,
			expected: `{"tools":[{"name":"my_tool"}]}`,
		},
		{
			name:     "no tools array returns body unchanged",
			forward:  map[string]string{"my_tool": "renamed"},
			body:     `{"messages":[]}`,
			expected: `{"messages":[]}`,
		},
		{
			name:     "tools is not array returns body unchanged",
			forward:  map[string]string{"my_tool": "renamed"},
			body:     `{"tools":"not_array"}`,
			expected: `{"tools":"not_array"}`,
		},
		{
			name:     "matching tool gets renamed",
			forward:  map[string]string{"my_tool": "renamed_tool"},
			body:     `{"tools":[{"name":"my_tool","description":"desc"}]}`,
			expected: `renamed_tool`,
		},
		{
			name:     "non-matching tool unchanged",
			forward:  map[string]string{"other_tool": "renamed"},
			body:     `{"tools":[{"name":"my_tool"}]}`,
			expected: `my_tool`,
		},
		{
			name:     "partial match - only exact match renames",
			forward:  map[string]string{"tool_a": "tool_x", "tool_b": "tool_y"},
			body:     `{"tools":[{"name":"tool_a"},{"name":"tool_c"},{"name":"tool_b"}]}`,
			expected: `tool_x`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Sanitizer{
				toolsForward: tt.forward,
				toolsReverse: make(map[string]string),
			}
			result := string(s.renameTools([]byte(tt.body)))
			if !strings.Contains(result, tt.expected) {
				t.Errorf("result %q does not contain %q", result, tt.expected)
			}
		})
	}
}

func TestRenameTools_MultipleTools(t *testing.T) {
	s := &Sanitizer{
		toolsForward: map[string]string{"tool_a": "tool_x", "tool_b": "tool_y"},
		toolsReverse: make(map[string]string),
	}
	body := `{"tools":[{"name":"tool_a"},{"name":"tool_c"},{"name":"tool_b"}]}`
	result := string(s.renameTools([]byte(body)))
	if !strings.Contains(result, `"tool_x"`) {
		t.Errorf("tool_a should be renamed to tool_x, got %s", result)
	}
	if !strings.Contains(result, `"tool_y"`) {
		t.Errorf("tool_b should be renamed to tool_y, got %s", result)
	}
	if !strings.Contains(result, `"tool_c"`) {
		t.Errorf("tool_c should remain unchanged, got %s", result)
	}
}

func TestReplaceSystem(t *testing.T) {
	tests := []struct {
		name     string
		rules    []config.ReplaceRule
		body     string
		contains string
	}{
		{
			name:     "empty rules returns body unchanged",
			rules:    nil,
			body:     `{"system":[{"type":"text","text":"hello world"}]}`,
			contains: "hello world",
		},
		{
			name:     "no system field returns body unchanged",
			rules:    []config.ReplaceRule{{Match: "hello", Replace: "goodbye"}},
			body:     `{"messages":[]}`,
			contains: `"messages":[]`,
		},
		{
			name:     "system not array returns body unchanged",
			rules:    []config.ReplaceRule{{Match: "hello", Replace: "goodbye"}},
			body:     `{"system":"just a string"}`,
			contains: "just a string",
		},
		{
			name:     "single block single rule",
			rules:    []config.ReplaceRule{{Match: "hello", Replace: "goodbye"}},
			body:     `{"system":[{"type":"text","text":"hello world"}]}`,
			contains: "goodbye world",
		},
		{
			name:     "multiple blocks",
			rules:    []config.ReplaceRule{{Match: "AAA", Replace: "BBB"}},
			body:     `{"system":[{"type":"text","text":"AAA first"},{"type":"text","text":"AAA second"}]}`,
			contains: "BBB first",
		},
		{
			name:     "multiple rules applied in order",
			rules:    []config.ReplaceRule{{Match: "cat", Replace: "dog"}, {Match: "dog", Replace: "fish"}},
			body:     `{"system":[{"type":"text","text":"I have a cat"}]}`,
			contains: "I have a fish",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Sanitizer{
				toolsForward: make(map[string]string),
				toolsReverse: make(map[string]string),
				systemRules:  tt.rules,
			}
			result := string(s.replaceSystem([]byte(tt.body)))
			if !strings.Contains(result, tt.contains) {
				t.Errorf("result %q does not contain %q", result, tt.contains)
			}
		})
	}
}

func TestReplaceSystem_MultipleBlocks(t *testing.T) {
	s := &Sanitizer{
		toolsForward: make(map[string]string),
		toolsReverse: make(map[string]string),
		systemRules:  []config.ReplaceRule{{Match: "AAA", Replace: "BBB"}},
	}
	body := `{"system":[{"type":"text","text":"AAA first"},{"type":"text","text":"AAA second"}]}`
	result := string(s.replaceSystem([]byte(body)))
	if !strings.Contains(result, "BBB first") {
		t.Errorf("first block not replaced: %s", result)
	}
	if !strings.Contains(result, "BBB second") {
		t.Errorf("second block not replaced: %s", result)
	}
}

func TestReplaceBody(t *testing.T) {
	tests := []struct {
		name     string
		rules    []config.ReplaceRule
		body     string
		expected string
	}{
		{
			name:     "empty rules returns body unchanged",
			rules:    nil,
			body:     `{"foo":"bar"}`,
			expected: `{"foo":"bar"}`,
		},
		{
			name:     "single replacement across entire body",
			rules:    []config.ReplaceRule{{Match: "SECRET", Replace: "REDACTED"}},
			body:     `{"data":"SECRET value SECRET"}`,
			expected: `{"data":"REDACTED value REDACTED"}`,
		},
		{
			name:     "multiple rules applied sequentially",
			rules:    []config.ReplaceRule{{Match: "AAA", Replace: "BBB"}, {Match: "BBB", Replace: "CCC"}},
			body:     `{"text":"AAA"}`,
			expected: `{"text":"CCC"}`,
		},
		{
			name:     "no match leaves body unchanged",
			rules:    []config.ReplaceRule{{Match: "NOMATCH", Replace: "X"}},
			body:     `{"text":"hello"}`,
			expected: `{"text":"hello"}`,
		},
		{
			name:     "empty body",
			rules:    []config.ReplaceRule{{Match: "a", Replace: "b"}},
			body:     ``,
			expected: ``,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Sanitizer{
				toolsForward: make(map[string]string),
				toolsReverse: make(map[string]string),
				bodyRules:    tt.rules,
			}
			result := string(s.replaceBody([]byte(tt.body)))
			if result != tt.expected {
				t.Errorf("got %q, want %q", result, tt.expected)
			}
		})
	}
}

func TestSanitizeRequest(t *testing.T) {
	cfg := config.SanitizeConfig{
		Tools:  []config.RenameRule{{From: "my_tool", To: "renamed_tool"}},
		System: []config.ReplaceRule{{Match: "INTERNAL", Replace: "PUBLIC"}},
		Body:   []config.ReplaceRule{{Match: "secret_val", Replace: "safe_val"}},
	}
	s := NewSanitizer(cfg)

	body := `{"tools":[{"name":"my_tool"}],"system":[{"type":"text","text":"INTERNAL info"}],"data":"secret_val here"}`
	result := string(s.SanitizeRequest([]byte(body)))

	if !strings.Contains(result, `"renamed_tool"`) {
		t.Errorf("tool not renamed in result: %s", result)
	}
	if !strings.Contains(result, "PUBLIC info") {
		t.Errorf("system not replaced in result: %s", result)
	}
	if !strings.Contains(result, "safe_val here") {
		t.Errorf("body not replaced in result: %s", result)
	}
	if strings.Contains(result, "secret_val") {
		t.Errorf("secret_val should have been replaced: %s", result)
	}
}

func TestSanitizeRequest_EmptyConfig(t *testing.T) {
	s := NewSanitizer(config.SanitizeConfig{})
	body := `{"tools":[{"name":"my_tool"}],"system":[{"type":"text","text":"hello"}]}`
	result := string(s.SanitizeRequest([]byte(body)))
	if result != body {
		t.Errorf("empty config should not modify body.\ngot:  %s\nwant: %s", result, body)
	}
}

func TestDesanitizeResponse(t *testing.T) {
	tests := []struct {
		name     string
		reverse  map[string]string
		body     string
		expected string
	}{
		{
			name:     "no content field returns unchanged",
			reverse:  map[string]string{"renamed": "original"},
			body:     `{"id":"msg_1","role":"assistant"}`,
			expected: `{"id":"msg_1","role":"assistant"}`,
		},
		{
			name:     "content not array returns unchanged",
			reverse:  map[string]string{"renamed": "original"},
			body:     `{"content":"just text"}`,
			expected: `{"content":"just text"}`,
		},
		{
			name:     "non-tool_use block left unchanged",
			reverse:  map[string]string{"renamed": "original"},
			body:     `{"content":[{"type":"text","text":"hello"}]}`,
			expected: `{"content":[{"type":"text","text":"hello"}]}`,
		},
		{
			name:     "tool_use block with matching name gets reversed",
			reverse:  map[string]string{"renamed_tool": "original_tool"},
			body:     `{"content":[{"type":"tool_use","name":"renamed_tool","id":"t1"}]}`,
			expected: `original_tool`,
		},
		{
			name:     "tool_use block with no match unchanged",
			reverse:  map[string]string{"other": "something"},
			body:     `{"content":[{"type":"tool_use","name":"my_tool","id":"t1"}]}`,
			expected: `my_tool`,
		},
		{
			name:     "mixed blocks only tool_use reversed",
			reverse:  map[string]string{"renamed": "original"},
			body:     `{"content":[{"type":"text","text":"hi"},{"type":"tool_use","name":"renamed","id":"t1"}]}`,
			expected: `original`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Sanitizer{
				toolsForward: make(map[string]string),
				toolsReverse: tt.reverse,
			}
			result := string(s.DesanitizeResponse([]byte(tt.body)))
			if !strings.Contains(result, tt.expected) {
				t.Errorf("result %q does not contain %q", result, tt.expected)
			}
		})
	}
}

func TestDesanitizeResponse_MultipleToolUse(t *testing.T) {
	s := &Sanitizer{
		toolsForward: make(map[string]string),
		toolsReverse: map[string]string{"r1": "o1", "r2": "o2"},
	}
	body := `{"content":[{"type":"tool_use","name":"r1","id":"t1"},{"type":"text","text":"x"},{"type":"tool_use","name":"r2","id":"t2"}]}`
	result := string(s.DesanitizeResponse([]byte(body)))
	if !strings.Contains(result, `"o1"`) {
		t.Errorf("r1 not reversed to o1: %s", result)
	}
	if !strings.Contains(result, `"o2"`) {
		t.Errorf("r2 not reversed to o2: %s", result)
	}
}

func TestDesanitizeStreamEvent(t *testing.T) {
	tests := []struct {
		name     string
		reverse  map[string]string
		line     string
		expected string
	}{
		{
			name:     "non-data line passed through",
			reverse:  map[string]string{"r": "o"},
			line:     "event: content_block_start",
			expected: "event: content_block_start",
		},
		{
			name:     "data line without tool_use passed through",
			reverse:  map[string]string{"r": "o"},
			line:     `data: {"type":"text","text":"hello"}`,
			expected: `data: {"type":"text","text":"hello"}`,
		},
		{
			name:     "data line with tool_use in content_block.name",
			reverse:  map[string]string{"renamed_tool": "original_tool"},
			line:     `data: {"type":"content_block_start","content_block":{"type":"tool_use","name":"renamed_tool","id":"t1"}}`,
			expected: `original_tool`,
		},
		{
			name:     "data line with tool_use in delta.name",
			reverse:  map[string]string{"renamed_tool": "original_tool"},
			line:     `data: {"type":"content_block_delta","delta":{"type":"tool_use","name":"renamed_tool"}}`,
			expected: `original_tool`,
		},
		{
			name:     "data line with tool_use but no matching name",
			reverse:  map[string]string{"other": "something"},
			line:     `data: {"type":"content_block_start","content_block":{"type":"tool_use","name":"my_tool","id":"t1"}}`,
			expected: `my_tool`,
		},
		{
			name:     "empty line passed through",
			reverse:  map[string]string{"r": "o"},
			line:     "",
			expected: "",
		},
		{
			name:     "line contains tool_use but not data prefix - passed through",
			reverse:  map[string]string{"r": "o"},
			line:     "event: tool_use",
			expected: "event: tool_use",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Sanitizer{
				toolsForward: make(map[string]string),
				toolsReverse: tt.reverse,
			}
			result := s.DesanitizeStreamEvent(tt.line)
			if !strings.Contains(result, tt.expected) {
				t.Errorf("result %q does not contain %q", result, tt.expected)
			}
		})
	}
}

func TestDesanitizeStreamEvent_DataPrefixPreserved(t *testing.T) {
	s := &Sanitizer{
		toolsForward: make(map[string]string),
		toolsReverse: map[string]string{"renamed": "original"},
	}
	line := `data: {"type":"content_block_start","content_block":{"type":"tool_use","name":"renamed","id":"t1"}}`
	result := s.DesanitizeStreamEvent(line)
	if !strings.HasPrefix(result, "data: ") {
		t.Errorf("result should start with 'data: ', got %q", result)
	}
}

func TestSanitizeRequest_MalformedJSON(t *testing.T) {
	s := NewSanitizer(config.SanitizeConfig{
		Tools:  []config.RenameRule{{From: "a", To: "b"}},
		System: []config.ReplaceRule{{Match: "x", Replace: "y"}},
	})
	// Malformed JSON - renameTools and replaceSystem should handle gracefully
	body := `not valid json`
	result := string(s.SanitizeRequest([]byte(body)))
	// Should not panic; body rules still do string replacement
	if result != "not valid json" {
		t.Errorf("malformed JSON should pass through (no body rules match), got %q", result)
	}
}

func TestSanitizeRequest_EmptyBody(t *testing.T) {
	s := NewSanitizer(config.SanitizeConfig{
		Tools: []config.RenameRule{{From: "a", To: "b"}},
	})
	result := s.SanitizeRequest([]byte{})
	if len(result) != 0 {
		t.Errorf("empty body should return empty, got %q", string(result))
	}
}