package usecase

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"testing"
	"unicode/utf8"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)

type recordingRepairClient struct {
	requests []domain.GenerateRequest
	response *domain.GenerateResponse
	err      error
}

func (c *recordingRepairClient) Generate(_ context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
	c.requests = append(c.requests, req)
	return c.response, c.err
}

func TestDefaultOutputRepairerBuildsFullContextRequest(t *testing.T) {
	client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
	repairer := NewDefaultOutputRepairer(client)
	original := []domain.RenderedMessage{
		{Role: "system", Content: "Follow the task.", CacheControl: &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"}},
		{Role: "user", Content: "Summarize the report."},
	}
	before := append([]domain.RenderedMessage(nil), original...)
	previous := strings.Repeat("candidate ", 12_000)
	target := domain.ExecutionTarget{BackendID: "backend", Endpoint: "https://provider.example/v1", Model: "model"}
	presence := domain.ExecutionTargetPresence{Temperature: true, TopP: true}
	structured := &domain.StructuredOutputSpec{}

	response, err := repairer.Repair(context.Background(), RepairRequest{
		OriginalMessages: original,
		PreviousOutput:   previous,
		ValidationErrors: []string{"invalid JSON"},
		SessionID:        "session",
		Target:           target,
		TargetPresence:   presence,
		StructuredOutput: structured,
		Attempt:          1,
		MaxAttempts:      3,
		Mode:             domain.ValidationJSONSchema,
	})
	if err != nil {
		t.Fatalf("repair: %v", err)
	}
	if response == nil || response.Content != "corrected" {
		t.Fatalf("response = %+v", response)
	}
	if !reflect.DeepEqual(original, before) {
		t.Fatalf("original messages changed: got %#v, want %#v", original, before)
	}
	if len(client.requests) != 1 {
		t.Fatalf("generation requests = %d, want 1", len(client.requests))
	}

	request := client.requests[0]
	if request.Prompt.SessionID != "session" || !reflect.DeepEqual(request.Target, target) || request.TargetPresence != presence || request.StructuredOutput != structured {
		t.Fatalf("generation request fields = %+v", request)
	}
	if len(request.Prompt.Messages) != len(original)+2 {
		t.Fatalf("message count = %d, want %d", len(request.Prompt.Messages), len(original)+2)
	}
	if !reflect.DeepEqual(request.Prompt.Messages[:len(original)], original) {
		t.Fatalf("original messages = %#v, want %#v", request.Prompt.Messages[:len(original)], original)
	}
	assistant := request.Prompt.Messages[len(original)]
	if assistant.Role != "assistant" || assistant.Content != previous {
		t.Fatalf("assistant candidate = %+v", assistant)
	}
	correction := request.Prompt.Messages[len(original)+1]
	if correction.Role != "user" || !strings.Contains(correction.Content, "Repair attempt 1 of 3") ||
		!strings.Contains(correction.Content, "Preserve valid values") ||
		!strings.Contains(correction.Content, "Return only corrected JSON") {
		t.Fatalf("correction message = %q", correction.Content)
	}
	if diagnostics := repairDiagnosticsFromMessage(t, correction.Content); !reflect.DeepEqual(diagnostics, []string{"invalid JSON"}) {
		t.Fatalf("diagnostics = %#v", diagnostics)
	}
}

func TestDefaultOutputRepairerDoesNotAccumulateCandidates(t *testing.T) {
	client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
	repairer := NewDefaultOutputRepairer(client)
	backing := make([]domain.RenderedMessage, 1, 4)
	backing[0] = domain.RenderedMessage{Role: "user", Content: "Original task"}
	before := append([]domain.RenderedMessage(nil), backing...)

	for _, candidate := range []string{"first invalid", "second invalid"} {
		_, err := repairer.Repair(context.Background(), RepairRequest{
			OriginalMessages: backing,
			PreviousOutput:   candidate,
			Attempt:          1,
			MaxAttempts:      3,
			Mode:             domain.ValidationJSON,
		})
		if err != nil {
			t.Fatalf("repair %q: %v", candidate, err)
		}
	}
	if !reflect.DeepEqual(backing, before) {
		t.Fatalf("caller messages changed: got %#v, want %#v", backing, before)
	}
	if len(client.requests) != 2 {
		t.Fatalf("generation requests = %d, want 2", len(client.requests))
	}
	for index, request := range client.requests {
		messages := request.Prompt.Messages
		if len(messages) != 3 || messages[0] != backing[0] || messages[1].Role != "assistant" || messages[1].Content != []string{"first invalid", "second invalid"}[index] {
			t.Fatalf("request %d messages = %#v", index, messages)
		}
	}
}

func TestDefaultOutputRepairerOmitsEmptyCandidateMessage(t *testing.T) {
	for _, tc := range []struct {
		name      string
		candidate string
	}{
		{name: "empty", candidate: ""},
		{name: "whitespace", candidate: " \n\t "},
	} {
		t.Run(tc.name, func(t *testing.T) {
			client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
			repairer := NewDefaultOutputRepairer(client)
			_, err := repairer.Repair(context.Background(), RepairRequest{
				OriginalMessages: []domain.RenderedMessage{{Role: "user", Content: "Original task"}},
				PreviousOutput:   tc.candidate,
				Attempt:          1,
				MaxAttempts:      1,
				Mode:             domain.ValidationBasic,
			})
			if err != nil {
				t.Fatalf("repair: %v", err)
			}
			messages := client.requests[0].Prompt.Messages
			if len(messages) != 2 || messages[1].Role != "user" || !strings.Contains(messages[1].Content, "previous response was empty") {
				t.Fatalf("messages = %#v", messages)
			}
			if strings.Contains(messages[1].Content, tc.candidate) && tc.candidate != "" {
				t.Fatalf("correction message repeated whitespace candidate: %q", messages[1].Content)
			}
		})
	}
}

func TestDefaultOutputRepairerUsesModeSpecificGuidance(t *testing.T) {
	tests := []struct {
		mode domain.ValidationMode
		want string
	}{
		{mode: domain.ValidationBasic, want: "Return a nonempty response"},
		{mode: domain.ValidationJSON, want: "Return only corrected JSON"},
		{mode: domain.ValidationJSONSchema, want: "Return only corrected JSON"},
	}

	for _, tc := range tests {
		t.Run(string(tc.mode), func(t *testing.T) {
			client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
			_, err := NewDefaultOutputRepairer(client).Repair(context.Background(), RepairRequest{Mode: tc.mode, Attempt: 1, MaxAttempts: 1})
			if err != nil {
				t.Fatalf("repair: %v", err)
			}
			if !strings.Contains(client.requests[0].Prompt.Messages[0].Content, tc.want) {
				t.Fatalf("correction message = %q", client.requests[0].Prompt.Messages[0].Content)
			}
		})
	}

	client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "unexpected"}}
	_, err := NewDefaultOutputRepairer(client).Repair(context.Background(), RepairRequest{Mode: domain.ValidationNone})
	if err == nil || !strings.Contains(err.Error(), "unsupported validation mode") {
		t.Fatalf("unsupported mode error = %v", err)
	}
	if len(client.requests) != 0 {
		t.Fatalf("generation requests = %d, want 0", len(client.requests))
	}
}

func TestFormatRepairDiagnosticsBoundsAndPreservesData(t *testing.T) {
	tests := []struct {
		name         string
		errors       []string
		wantOmission bool
		check        func(*testing.T, []string)
	}{
		{
			name:   "below limit",
			errors: []string{"first", "second"},
			check: func(t *testing.T, got []string) {
				t.Helper()
				if !reflect.DeepEqual(got, []string{"first", "second"}) {
					t.Fatalf("diagnostics = %#v", got)
				}
			},
		},
		{
			name:   "at limit",
			errors: []string{strings.Repeat("x", maxRepairDiagnosticBytes-4)},
			check: func(t *testing.T, got []string) {
				t.Helper()
				if len(got) != 1 || len(got[0]) != maxRepairDiagnosticBytes-4 {
					t.Fatalf("diagnostics lengths = %#v", got)
				}
			},
		},
		{
			name:         "multibyte truncation preserves prior entries",
			errors:       []string{"first", strings.Repeat("界", maxRepairDiagnosticBytes)},
			wantOmission: true,
			check: func(t *testing.T, got []string) {
				t.Helper()
				if len(got) != 3 || got[0] != "first" || !strings.HasPrefix(strings.Repeat("界", maxRepairDiagnosticBytes), got[1]) {
					t.Fatalf("diagnostics = %#v", got)
				}
			},
		},
		{
			name:         "many diagnostics",
			errors:       manyRepairDiagnostics(),
			wantOmission: true,
			check: func(t *testing.T, got []string) {
				t.Helper()
				if len(got) < 2 || got[0] != manyRepairDiagnostics()[0] {
					t.Fatalf("diagnostics = %#v", got)
				}
			},
		},
		{
			name: "no room for partial diagnostic",
			errors: func() []string {
				omission, _ := json.Marshal(omittedRepairDiagnostics)
				return []string{
					strings.Repeat("x", maxRepairDiagnosticBytes-len(omission)-5),
					strings.Repeat("y", 128),
				}
			}(),
			wantOmission: true,
			check: func(t *testing.T, got []string) {
				t.Helper()
				if len(got) != 2 || got[1] != omittedRepairDiagnostics {
					t.Fatalf("diagnostics = %#v", got)
				}
			},
		},
		{
			name:   "invalid UTF-8",
			errors: []string{"broken\xffinput"},
			check: func(t *testing.T, got []string) {
				t.Helper()
				if !reflect.DeepEqual(got, []string{"broken\uFFFDinput"}) {
					t.Fatalf("diagnostics = %#v", got)
				}
			},
		},
		{
			name:         "one huge diagnostic",
			errors:       []string{strings.Repeat("x", maxRepairDiagnosticBytes*2)},
			wantOmission: true,
			check: func(t *testing.T, got []string) {
				t.Helper()
				if len(got) != 2 || !strings.HasPrefix(strings.Repeat("x", maxRepairDiagnosticBytes*2), got[0]) {
					t.Fatalf("diagnostics = %#v", got)
				}
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			before := append([]string(nil), tc.errors...)
			encoded := formatRepairDiagnostics(tc.errors)
			if len(encoded) > maxRepairDiagnosticBytes || !utf8.ValidString(encoded) {
				t.Fatalf("encoded diagnostic length/UTF-8 = (%d, %v)", len(encoded), utf8.ValidString(encoded))
			}
			var got []string
			if err := json.Unmarshal([]byte(encoded), &got); err != nil {
				t.Fatalf("decode diagnostics: %v; encoded=%q", err, encoded)
			}
			if !reflect.DeepEqual(tc.errors, before) {
				t.Fatalf("input diagnostics changed: got %#v, want %#v", tc.errors, before)
			}
			if hasOmission := len(got) > 0 && got[len(got)-1] == omittedRepairDiagnostics; hasOmission != tc.wantOmission {
				t.Fatalf("omission = %v, want %v; diagnostics=%#v", hasOmission, tc.wantOmission, got)
			}
			tc.check(t, got)
		})
	}
}

func manyRepairDiagnostics() []string {
	diagnostics := make([]string, 1_000)
	for index := range diagnostics {
		diagnostics[index] = fmt.Sprintf("diagnostic %04d %s", index, strings.Repeat("x", 128))
	}
	return diagnostics
}

func TestDefaultOutputRepairerPropagatesGenerationFailures(t *testing.T) {
	expected := errors.New("generation failed")
	for _, tc := range []struct {
		name   string
		client *recordingRepairClient
		want   error
	}{
		{name: "nil client", want: nil},
		{name: "generation error", client: &recordingRepairClient{err: expected}, want: expected},
		{name: "nil response", client: &recordingRepairClient{}, want: nil},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var repairer OutputRepairer
			if tc.client == nil {
				repairer = NewDefaultOutputRepairer(nil)
			} else {
				repairer = NewDefaultOutputRepairer(tc.client)
			}
			response, err := repairer.Repair(context.Background(), RepairRequest{Mode: domain.ValidationJSON})
			if response != nil || err == nil {
				t.Fatalf("response/error = (%+v, %v)", response, err)
			}
			if tc.want != nil && !errors.Is(err, tc.want) {
				t.Fatalf("error = %v, want %v", err, tc.want)
			}
		})
	}
}

func repairDiagnosticsFromMessage(t *testing.T, message string) []string {
	t.Helper()
	const marker = "Validation diagnostics (data):\n"
	index := strings.Index(message, marker)
	if index < 0 {
		t.Fatalf("missing diagnostics marker in %q", message)
	}
	encoded := message[index+len(marker):]
	var diagnostics []string
	if err := json.Unmarshal([]byte(encoded), &diagnostics); err != nil {
		t.Fatalf("decode diagnostics: %v", err)
	}
	return diagnostics
}