package processreport

import (
	"context"
	"errors"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/audita/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/internal/core/schema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/audita/internal/framework/outcomes"
	"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/internal/framework/runner"
	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/internal/framework/validators"
)

type partialModule struct{ validator contracts.Validator }

func (m partialModule) Key() string { return "grammar" }
func (m partialModule) ReplacementPolicy() proposals.ReplacementPolicy {
	return proposals.ReplacementPolicyRequireUnique
}
func (m partialModule) Validators() []contracts.Validator { return []contracts.Validator{m.validator} }
func (m partialModule) ModuleForSpec(contracts.ModuleRunSpec) (contracts.TranscriptModule, error) {
	return m, nil
}
func (m partialModule) Propose(context.Context, contracts.ProposalRequest) (contracts.ProposalResult, error) {
	result := contracts.ProposalResult{}
	for _, text := range []string{"Hello", "Greetings", "Hi"} {
		result.Proposals = append(result.Proposals, proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "hello", CorrectedText: text, Confidence: 1})
	}
	return result, nil
}

type oneItemEstimator struct{}

func (oneItemEstimator) EstimateTokens(payload string) int {
	return 100 * strings.Count(payload, `"correction_index"`)
}

type interruptedValidationClient struct {
	cause error
	calls int
}

func (c *interruptedValidationClient) CompleteStructured(_ context.Context, req structuredoutput.Request, target any) (structuredoutput.Response, error) {
	c.calls++
	if *req.Context.BatchIndex == 2 {
		return structuredoutput.Response{}, &structuredoutput.Failure{Class: structuredoutput.FailureFatal, Reason: outcomes.ReasonSetupFailure, Message: "validation interrupted", Cause: c.cause}
	}
	index := *req.Context.BatchIndex
	if index == 0 {
		*target.(*validators.LLMValidationResponse) = validators.LLMValidationResponse{}
		return structuredoutput.Response{}, nil
	}
	*target.(*validators.LLMValidationResponse) = validators.LLMValidationResponse{Validations: []validators.LLMValidationDecision{{CorrectionIndex: index, Approved: index == 1, Confidence: 1, Reason: "completed assessment"}}}
	return structuredoutput.Response{}, nil
}

func TestFatalBatchRetainsEarlierDecisionsInLedger(t *testing.T) {
	validator, err := validators.NewLLMBackedValidator("editorial_review", validators.LLMValidatorTypeEditorialReview)
	if err != nil {
		t.Fatal(err)
	}
	validator.SetTokenEstimator(oneItemEstimator{})
	origin := errors.New("later batch failed")
	client := &interruptedValidationClient{cause: origin}
	cfg := config.Config{Chunking: config.Chunking{MaxSectionTokens: 8192, MinSectionTokens: 1}, Validation: config.Validation{MaxPromptTokens: 150}}
	result, err := runner.New(partialModule{validator}).Run(context.Background(), runner.RunInput{
		Config: &cfg, Transcript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "hello"}}},
		ModuleSpecs: []contracts.ModuleRunSpec{{ModuleKey: "grammar", InstanceName: "grammar"}}, ValidationLLMClient: client, CaptureCandidates: true,
	})
	if !errors.Is(err, origin) || client.calls != 3 {
		t.Fatalf("fatal cause/calls: %v %d", err, client.calls)
	}
	if result.FinalTranscript.Segments[0].Text != "hello" || len(result.ModuleResults[0].AppliedChanges) != 0 {
		t.Fatal("fatal partial approvals changed transcript")
	}
	if len(result.WorkUnits[0].Operations) != 3 || result.WorkUnits[0].State != outcomes.WorkInterrupted {
		t.Fatalf("partial facts: %+v", result.WorkUnits)
	}
	ledger := BuildCorrectionLedger(CorrectionLedgerInput{RunOutput: &result})
	candidates := make(map[int]CorrectionLedgerEntry)
	for _, entry := range ledger {
		if entry.Disposition != correctionDispositionFailed {
			candidates[entry.ProposalIndex] = entry
		}
	}
	if len(candidates) != 3 {
		t.Fatalf("candidate evidence: %+v", ledger)
	}
	for index := 0; index < 2; index++ {
		entry := candidates[index]
		wantDisposition := validators.DispositionUnavailable
		if index == 1 {
			wantDisposition = validators.DispositionApproved
		}
		if len(entry.LLMValidatorResults) != 1 || entry.LLMValidatorResults[0].Approved != (index == 1) || entry.LLMValidatorResults[0].Disposition != wantDisposition {
			t.Errorf("lost batch %d decision: %+v", index, entry)
		}
	}
	if candidates[0].Disposition != correctionDispositionUnavailable || candidates[1].Disposition != correctionDispositionInterrupted || candidates[2].Disposition != correctionDispositionInterrupted {
		t.Fatalf("wrong partial dispositions: %+v", ledger)
	}
}