package npcregistry

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-alpha:chunk:0",
		SourceID:  doc.ID,
		Index:     0,
		Ref:       source.SourceRef{SourceID: doc.ID, StartUnitID: 1, EndUnitID: 3},
		Content:   []byte(`{"units":[1,2,3]}`),
		MediaType: "application/json",
		Units:     append([]source.SourceUnit(nil), doc.Units...),
		Metadata:  map[string]any{"ignored": "chunk metadata"},
	}
	return contracts.TypedExtractionRequest{
		Source:      doc,
		Chunk:       chunk,
		SourceInput: contracts.NewLLMInputMaterial("source", chunk.MediaType, chunk.Content, "sha256:chunk", "file:///session-alpha.json"),
		SessionID:   "session-123",
		LLMProfile:  "profile-npcs",
	}
}

func sourceDocument() *source.SourceDocument {
	return &source.SourceDocument{
		ID:     "session-alpha",
		Kind:   "transcript",
		Format: "application/vnd.seriatim.minimal+json",
		Digest: "sha256:test",
		Units: []source.SourceUnit{
			{ID: 1, Kind: "transcript_segment", Text: "Mira Thorn greets the party."},
			{ID: 2, Kind: "transcript_segment", Text: "The Greencloak watches the northern road."},
			{ID: 3, Kind: "transcript_segment", Text: "Captain Vale reports to Mira Thorn."},
		},
	}
}

func responseSourceRefs(startUnitID, endUnitID int) []npcSourceRefResponse {
	return []npcSourceRefResponse{{StartUnitID: startUnitID, EndUnitID: endUnitID}}
}

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, want nil", err)
	}
	return extractor
}

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

func cloneStructuredCompletionRequest(req contracts.StructuredCompletionRequest) contracts.StructuredCompletionRequest {
	req.Inputs = req.Inputs.Clone()
	correction, err := contracts.CloneSemanticCorrection(req.Correction)
	if err != nil {
		panic(err)
	}
	req.Correction = correction
	if len(req.Vars) == 0 {
		req.Vars = nil
		return req
	}
	vars := make(map[string]any, len(req.Vars))
	for key, value := range req.Vars {
		vars[key] = value
	}
	req.Vars = vars
	return req
}

type fakeNPCsLLMClient struct {
	response extractionResponse
	content  []byte
	err      error
	requests []contracts.StructuredCompletionRequest
}

func (client *fakeNPCsLLMClient) 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 := append([]byte(nil), client.content...)
	if len(content) != 0 {
		if err := json.Unmarshal(content, target); err != nil {
			return contracts.StructuredCompletionResponse{}, err
		}
	} else {
		*target = client.response
		var err error
		content, err = json.Marshal(client.response)
		if err != nil {
			return contracts.StructuredCompletionResponse{}, err
		}
	}
	return contracts.StructuredCompletionResponse{Content: content}, nil
}
