package locationoccurrences

import (
	"context"
	"encoding/json"
	"errors"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

func extractionRequest() contracts.TypedExtractionRequest {
	doc := sourceDocument()
	chunk := &source.Chunk{
		ID: "session-occurrences:chunk:0", SourceID: doc.ID, Index: 0,
		Ref:     source.SourceRef{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 30},
		Content: []byte(`{"units":[10,20,30]}`), MediaType: "application/json",
		Units: append([]source.SourceUnit(nil), doc.Units...),
	}
	return contracts.TypedExtractionRequest{
		Source: doc, Chunk: chunk,
		SourceInput: contracts.NewLLMInputMaterial("source", chunk.MediaType, chunk.Content, "sha256:chunk", "file:///session-occurrences.json"),
		SessionID:   "occurrence-session", LLMProfile: "occurrence-profile",
	}
}

func sourceDocument() *source.SourceDocument {
	return &source.SourceDocument{ID: "session-occurrences", Kind: "transcript", Format: "application/json", Digest: "sha256:test", Units: []source.SourceUnit{
		{ID: 10, Kind: "transcript_segment", Text: "The party returns to the tavern."},
		{ID: 20, Kind: "transcript_segment", Text: "They plan to travel to the tavern tomorrow."},
		{ID: 30, Kind: "transcript_segment", Text: "They recall their first visit to the tavern."},
	}}
}

func occurrenceRefs(start, end int) []occurrenceSourceRefResponse {
	return []occurrenceSourceRefResponse{{StartUnitID: start, EndUnitID: end}}
}

func newExtractor(t *testing.T, client contracts.StructuredLLMClient, references ...contracts.ReferenceSet) *Extractor {
	t.Helper()
	extractor, err := New(client, Options{}, references...)
	if err != nil {
		t.Fatalf("New() error = %v", err)
	}
	return extractor
}

func mismatchedSourceInputRequest(req contracts.TypedExtractionRequest) contracts.TypedExtractionRequest {
	req.SourceInput = contracts.NewLLMInputMaterial("source", "application/json", []byte(`{"different":true}`), "sha256:other", "")
	return req
}

type fakeOccurrencesLLMClient struct {
	response extractionResponse
	err      error
	requests []contracts.StructuredCompletionRequest
}

func (client *fakeOccurrencesLLMClient) CompleteStructured(_ context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	client.requests = append(client.requests, cloneStructuredCompletionRequest(req))
	if client.err != nil {
		return contracts.StructuredCompletionResponse{}, client.err
	}
	target, ok := out.(*extractionResponse)
	if !ok {
		return contracts.StructuredCompletionResponse{}, errors.New("unexpected output target")
	}
	content, err := json.Marshal(client.response)
	if err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	if err := json.Unmarshal(content, target); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	return contracts.StructuredCompletionResponse{Content: content}, nil
}

func cloneStructuredCompletionRequest(req contracts.StructuredCompletionRequest) contracts.StructuredCompletionRequest {
	req.Inputs = req.Inputs.Clone()
	return req
}