package proposal_generation

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/audita/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/audita/internal/framework/llm"
	"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
)

type fakeStructuredClient struct {
	responses []StructuredCorrectionSet
	err       error
	calls     []contracts.StructuredCompletionRequest
}

func (f *fakeStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	_ = ctx
	f.calls = append(f.calls, req)
	if f.err != nil {
		return contracts.StructuredCompletionResponse{}, f.err
	}
	target, ok := out.(*StructuredCorrectionSet)
	if !ok {
		return contracts.StructuredCompletionResponse{}, errors.New("unexpected output type")
	}
	if len(f.responses) == 0 {
		return contracts.StructuredCompletionResponse{}, errors.New("unexpected call")
	}
	*target = f.responses[0]
	f.responses = f.responses[1:]
	return contracts.StructuredCompletionResponse{}, nil
}

type countingScheduler struct {
	runs int
}

func (s *countingScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
	s.runs++
	return fn(ctx)
}

type captureDiagnosticsWriter struct {
	lastStage           string
	lastRequestMetadata any
	lastRequestPayload  any
}

func (w *captureDiagnosticsWriter) WriteInteraction(stage string, requestMetadata any, requestPayload any, responsePayload any, errorPayload any) (InteractionArtifacts, error) {
	w.lastStage = stage
	w.lastRequestMetadata = requestMetadata
	w.lastRequestPayload = requestPayload
	_ = responsePayload
	_ = errorPayload
	return InteractionArtifacts{}, nil
}

type sleepingStructuredClient struct {
	inFlight    int32
	maxInFlight int32
	entered     chan struct{}
	release     chan struct{}
}

func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	_ = req
	current := atomic.AddInt32(&c.inFlight, 1)
	for {
		prior := atomic.LoadInt32(&c.maxInFlight)
		if current <= prior || atomic.CompareAndSwapInt32(&c.maxInFlight, prior, current) {
			break
		}
	}
	if c.entered != nil {
		c.entered <- struct{}{}
	}
	select {
	case <-c.release:
	case <-ctx.Done():
		atomic.AddInt32(&c.inFlight, -1)
		return contracts.StructuredCompletionResponse{}, ctx.Err()
	}
	atomic.AddInt32(&c.inFlight, -1)

	target, ok := out.(*StructuredCorrectionSet)
	if !ok {
		return contracts.StructuredCompletionResponse{}, errors.New("unexpected output type")
	}
	*target = StructuredCorrectionSet{
		Corrections: []StructuredCorrectionProposal{
			{TargetSegmentID: 1, OriginalText: "x", CorrectedText: "y", Confidence: 0.9},
		},
	}
	return contracts.StructuredCompletionResponse{}, nil
}

func defaultRequest(t *testing.T) Request {
	t.Helper()
	cfg := config.Default()
	return Request{
		ModuleKey:         "test_module",
		ModuleInstance:    "test_module",
		ReplacementPolicy: proposals.ReplacementPolicyRequireUnique,
		Config:            &cfg,
		Messages: []contracts.LLMMessage{
			{Role: "system", Content: "system prompt"},
			{Role: "user", Content: "user prompt"},
		},
		StartIndex: 0,
	}
}

func TestGenerateCandidatesSuccess(t *testing.T) {
	client := &fakeStructuredClient{
		responses: []StructuredCorrectionSet{
			{
				Corrections: []StructuredCorrectionProposal{
					{TargetSegmentID: 7, OriginalText: "teh", CorrectedText: "the", Confidence: 0.9},
				},
			},
		},
	}
	section := contracts.SectionMetadata{Index: 2}
	req := defaultRequest(t)
	req.LLMClient = client
	req.StartIndex = 10
	req.Section = &section

	got, err := GenerateCandidates(context.Background(), req)
	if err != nil {
		t.Fatalf("GenerateCandidates error: %v", err)
	}
	if len(got.Corrections) != 1 || len(got.Enriched) != 1 {
		t.Fatalf("unexpected proposal lengths: %+v", got)
	}
	if got.Enriched[0].ProposalIndex != 10 {
		t.Fatalf("expected proposal index 10, got %d", got.Enriched[0].ProposalIndex)
	}
	if got.Enriched[0].SectionIndex == nil || *got.Enriched[0].SectionIndex != 2 {
		t.Fatalf("expected section index 2, got %v", got.Enriched[0].SectionIndex)
	}
}

func TestGenerateCandidatesUsesCorrectionSetSchema(t *testing.T) {
	client := &fakeStructuredClient{
		responses: []StructuredCorrectionSet{
			{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "teh", CorrectedText: "the", Confidence: 0.9}}},
		},
	}
	req := defaultRequest(t)
	req.LLMClient = client

	_, err := GenerateCandidates(context.Background(), req)
	if err != nil {
		t.Fatalf("GenerateCandidates error: %v", err)
	}
	if len(client.calls) != 1 {
		t.Fatalf("expected 1 LLM call, got %d", len(client.calls))
	}
	call := client.calls[0]
	if call.ResponseSchema == nil {
		t.Fatalf("expected response schema on structured request")
	}
	want := responseschema.MustLookup(responseschema.CorrectionSetKey)
	if call.ResponseSchema.ID != want.ID || call.ResponseSchema.Version != want.Version || call.ResponseSchema.Name != want.Name || call.ResponseSchema.SHA256 != want.SHA256 {
		t.Fatalf("unexpected response schema metadata: got=%+v want=%+v", *call.ResponseSchema, want)
	}
}

func TestGenerateCandidatesDiagnosticsIncludeSchemaMetadata(t *testing.T) {
	client := &fakeStructuredClient{
		responses: []StructuredCorrectionSet{
			{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "teh", CorrectedText: "the", Confidence: 0.9}}},
		},
	}
	diag := &captureDiagnosticsWriter{}
	req := defaultRequest(t)
	req.LLMClient = client
	req.DiagnosticsWriter = diag

	_, err := GenerateCandidates(context.Background(), req)
	if err != nil {
		t.Fatalf("GenerateCandidates error: %v", err)
	}
	metadata, ok := diag.lastRequestMetadata.(map[string]any)
	if !ok {
		t.Fatalf("expected request metadata map, got %T", diag.lastRequestMetadata)
	}
	schemaMap, ok := metadata["response_schema"].(map[string]any)
	if !ok {
		t.Fatalf("expected response_schema metadata map, got %T", metadata["response_schema"])
	}
	want := responseschema.MustLookup(responseschema.CorrectionSetKey)
	if schemaMap["id"] != want.ID || schemaMap["version"] != want.Version || schemaMap["name"] != want.Name || schemaMap["sha256"] != want.SHA256 {
		t.Fatalf("unexpected diagnostics schema metadata: got=%v want=%+v", schemaMap, want)
	}
}

func TestGenerateCandidatesMalformedStructuredResponse(t *testing.T) {
	client := &fakeStructuredClient{
		responses: []StructuredCorrectionSet{
			{
				Corrections: []StructuredCorrectionProposal{
					{TargetSegmentID: 1, OriginalText: "x", CorrectedText: "", Confidence: 0.9},
				},
			},
		},
	}
	req := defaultRequest(t)
	req.LLMClient = client

	_, err := GenerateCandidates(context.Background(), req)
	if err == nil || !strings.Contains(err.Error(), "invalid structured correction") {
		t.Fatalf("expected structured response validation failure, got %v", err)
	}
}

func TestGenerateCandidatesDeterministicIndexAssignment(t *testing.T) {
	baseResponse := StructuredCorrectionSet{
		Corrections: []StructuredCorrectionProposal{
			{TargetSegmentID: 1, OriginalText: "a", CorrectedText: "A", Confidence: 0.9},
			{TargetSegmentID: 2, OriginalText: "b", CorrectedText: "B", Confidence: 0.9},
		},
	}
	clientA := &fakeStructuredClient{responses: []StructuredCorrectionSet{baseResponse}}
	clientB := &fakeStructuredClient{responses: []StructuredCorrectionSet{baseResponse}}

	reqA := defaultRequest(t)
	reqA.LLMClient = clientA
	reqA.StartIndex = 3
	first, err := GenerateCandidates(context.Background(), reqA)
	if err != nil {
		t.Fatalf("first generation failed: %v", err)
	}

	reqB := defaultRequest(t)
	reqB.LLMClient = clientB
	reqB.StartIndex = 3
	second, err := GenerateCandidates(context.Background(), reqB)
	if err != nil {
		t.Fatalf("second generation failed: %v", err)
	}

	if !reflect.DeepEqual(first.Enriched, second.Enriched) {
		t.Fatalf("expected stable enriched proposals\nfirst=%+v\nsecond=%+v", first.Enriched, second.Enriched)
	}
}

func TestGenerateCandidatesMultipleSectionsStableMetadata(t *testing.T) {
	client := &fakeStructuredClient{
		responses: []StructuredCorrectionSet{
			{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "alpha", CorrectedText: "ALPHA", Confidence: 0.9}}},
			{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 3, OriginalText: "charlie", CorrectedText: "CHARLIE", Confidence: 0.9}}},
		},
	}
	section0 := contracts.SectionMetadata{Index: 0}
	section1 := contracts.SectionMetadata{Index: 1}

	req0 := defaultRequest(t)
	req0.LLMClient = client
	req0.Section = &section0
	req0.StartIndex = 0
	part0, err := GenerateCandidates(context.Background(), req0)
	if err != nil {
		t.Fatalf("section 0 generation failed: %v", err)
	}

	req1 := defaultRequest(t)
	req1.LLMClient = client
	req1.Section = &section1
	req1.StartIndex = len(part0.Enriched)
	part1, err := GenerateCandidates(context.Background(), req1)
	if err != nil {
		t.Fatalf("section 1 generation failed: %v", err)
	}

	all := append(append([]proposals.EnrichedCorrectionProposal(nil), part0.Enriched...), part1.Enriched...)
	if len(all) != 2 {
		t.Fatalf("expected 2 proposals, got %d", len(all))
	}
	if all[0].ProposalIndex != 0 || all[1].ProposalIndex != 1 {
		t.Fatalf("unexpected proposal indexes: %d, %d", all[0].ProposalIndex, all[1].ProposalIndex)
	}
	if all[0].SectionIndex == nil || *all[0].SectionIndex != 0 || all[1].SectionIndex == nil || *all[1].SectionIndex != 1 {
		t.Fatalf("unexpected section metadata: %+v", all)
	}
}

func TestGenerateCandidatesDiagnosticsWrittenAndRedacted(t *testing.T) {
	secret := "proposal-secret"
	client := &fakeStructuredClient{
		responses: []StructuredCorrectionSet{
			{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: secret, CorrectedText: "safe", Confidence: 0.9}}},
		},
	}
	cfg := config.Default()
	cfg.PrimaryLLM.APIKey = secret
	req := defaultRequest(t)
	req.Config = &cfg
	req.LLMClient = client
	req.DiagnosticsDir = t.TempDir()
	req.Messages = []contracts.LLMMessage{
		{Role: "system", Content: "include secret " + secret},
		{Role: "user", Content: "fix it"},
	}

	got, err := GenerateCandidates(context.Background(), req)
	if err != nil {
		t.Fatalf("GenerateCandidates error: %v", err)
	}
	if got.Artifacts.ResponsePayloadPath == "" || got.Artifacts.RequestPayloadPath == "" {
		t.Fatalf("expected diagnostics artifact paths, got %+v", got.Artifacts)
	}

	for _, path := range []string{got.Artifacts.RequestPayloadPath, got.Artifacts.ResponsePayloadPath} {
		raw, readErr := os.ReadFile(path)
		if readErr != nil {
			t.Fatalf("read artifact %q: %v", path, readErr)
		}
		if strings.Contains(string(raw), secret) {
			t.Fatalf("artifact leaked secret %q: %s", path, string(raw))
		}
		if !strings.Contains(string(raw), "[REDACTED]") {
			t.Fatalf("expected redaction marker in artifact %q: %s", path, string(raw))
		}
	}

	raw, readErr := os.ReadFile(got.Artifacts.RequestMetadataPath)
	if readErr != nil {
		t.Fatalf("read metadata artifact %q: %v", got.Artifacts.RequestMetadataPath, readErr)
	}
	var metadata map[string]any
	if err := json.Unmarshal(raw, &metadata); err != nil {
		t.Fatalf("unmarshal metadata artifact: %v", err)
	}
	schemaMap, ok := metadata["response_schema"].(map[string]any)
	if !ok {
		t.Fatalf("expected response_schema metadata in diagnostics, got %T", metadata["response_schema"])
	}
	want := responseschema.MustLookup(responseschema.CorrectionSetKey)
	if schemaMap["id"] != want.ID || schemaMap["version"] != want.Version || schemaMap["name"] != want.Name || schemaMap["sha256"] != want.SHA256 {
		t.Fatalf("unexpected schema metadata in diagnostics: got=%v want=%+v", schemaMap, want)
	}
}

func TestGenerateCandidatesSchedulerUsage(t *testing.T) {
	client := &fakeStructuredClient{
		responses: []StructuredCorrectionSet{
			{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "x", CorrectedText: "y", Confidence: 0.9}}},
		},
	}
	scheduler := &countingScheduler{}
	req := defaultRequest(t)
	req.LLMClient = client
	req.Scheduler = scheduler

	_, err := GenerateCandidates(context.Background(), req)
	if err != nil {
		t.Fatalf("GenerateCandidates error: %v", err)
	}
	if scheduler.runs != 1 {
		t.Fatalf("expected scheduler to run once, got %d", scheduler.runs)
	}
}

func TestGenerateCandidatesClientError(t *testing.T) {
	client := &fakeStructuredClient{err: errors.New("boom")}
	req := defaultRequest(t)
	req.LLMClient = client
	req.DiagnosticsDir = t.TempDir()

	got, err := GenerateCandidates(context.Background(), req)
	if err == nil || !strings.Contains(err.Error(), "completion failed") {
		t.Fatalf("expected completion failure, got %v", err)
	}
	if got.Artifacts.ResponsePayloadPath != "" {
		t.Fatalf("expected zero result on error, got %+v", got)
	}
	matches, globErr := filepath.Glob(filepath.Join(req.DiagnosticsDir, req.ModuleInstance, "*error-payload.json"))
	if globErr != nil {
		t.Fatalf("glob error: %v", globErr)
	}
	if len(matches) == 0 {
		t.Fatalf("expected error diagnostics artifact under %s", req.DiagnosticsDir)
	}
}

func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
	scheduler, err := llm.NewScheduler(2)
	if err != nil {
		t.Fatalf("NewScheduler: %v", err)
	}
	client := &sleepingStructuredClient{
		entered: make(chan struct{}, 16),
		release: make(chan struct{}),
	}
	baseReq := defaultRequest(t)
	baseReq.LLMClient = client
	baseReq.Scheduler = scheduler

	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			req := baseReq
			req.StartIndex = idx
			if _, runErr := GenerateCandidates(context.Background(), req); runErr != nil {
				t.Errorf("GenerateCandidates[%d] error: %v", idx, runErr)
			}
		}(i)
	}

	waitForEntries(t, client.entered, 2)
	close(client.release)
	wg.Wait()

	if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
		t.Fatalf("expected scheduler cap <= 2, got %d", got)
	}
	if got := atomic.LoadInt32(&client.maxInFlight); got < 2 {
		t.Fatalf("expected observed concurrency of at least 2, got %d", got)
	}
}

func waitForEntries(t *testing.T, entered <-chan struct{}, want int) {
	t.Helper()
	deadline := time.Now().Add(300 * time.Millisecond)
	got := 0
	for got < want && time.Now().Before(deadline) {
		select {
		case <-entered:
			got++
		default:
			runtime.Gosched()
		}
	}
	if got < want {
		t.Fatalf("timed out waiting for %d entered calls (got %d)", want, got)
	}
}