package integration_test

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
	"sync"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/config"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/evidencecontext"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	locationoccurrencecodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locationoccurrences"
	locationcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locationregistry"
	locationoccurrences "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/locationoccurrences"
	locations "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/locationregistry"
	locationnormalize "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/normalize/locationregistry"
)

func TestLocationRegistryHandoffProducesOccurrencesAndEvidence(t *testing.T) {
	registries := productionNPCRegistries(t)
	resolved := resolveLocationPipeline(t, registries)
	client := &locationHandoffLLMClient{}
	output, err := runPreparedPipeline(t, registries, resolved, client, pipeline.RunInput{RawInput: readNPCFixture(t)})
	if err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	if len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 2 {
		t.Fatalf("run outputs = %#v rejected = %#v, want accepted registry and occurrences", output.NormalizeOutputs, output.Rejected)
	}

	registryOutput := normalizedLane(t, output, "locations")
	if registryOutput.Artifact.Kind != dnd.LocationRegistryKind || registryOutput.Artifact.Schema.ID != locationcodec.SchemaID {
		t.Fatalf("registry envelope = %#v, want durable location artifact", registryOutput)
	}
	registry, err := locationcodec.New().Decode(registryOutput.Artifact.Content)
	if err != nil {
		t.Fatalf("Decode(registry) error = %v", err)
	}
	if len(registry.Locations) != 2 || registry.Locations[0].Name != "Moon Gate" || registry.Locations[1].Name != "Moon Gate" || registry.Locations[0].ID == registry.Locations[1].ID {
		t.Fatalf("registry = %#v, want same-name locations distinguished by source-derived IDs", registry)
	}

	occurrenceOutput := normalizedLane(t, output, "occurrences")
	if occurrenceOutput.Artifact.Kind != dnd.LocationOccurrenceListKind || occurrenceOutput.Artifact.Schema.ID != locationoccurrencecodec.SchemaID {
		t.Fatalf("occurrence envelope = %#v, want durable occurrence artifact", occurrenceOutput)
	}
	occurrences, err := locationoccurrencecodec.New().Decode(occurrenceOutput.Artifact.Content)
	if err != nil {
		t.Fatalf("Decode(occurrences) error = %v", err)
	}
	if len(occurrences.Occurrences) != 2 || occurrences.Occurrences[0].LocationID == occurrences.Occurrences[1].LocationID || occurrences.Occurrences[0].LocationID != registry.Locations[0].ID || occurrences.Occurrences[1].LocationID != registry.Locations[1].ID {
		t.Fatalf("occurrences = %#v, want exact registry ID grounding", occurrences)
	}
	for _, occurrence := range occurrences.Occurrences {
		if occurrence.SourceRefs[0].SourceID != "npc-session" {
			t.Fatalf("occurrence evidence = %#v, want current source evidence only", occurrence.SourceRefs)
		}
	}

	request := client.requestFor(t, locationoccurrences.PromptID)
	registryInput := request.Inputs["location_registry"]
	if registryInput.MediaType != locationcodec.MediaType || strings.Contains(string(registryInput.Content), "source_refs") || strings.Contains(string(registryInput.Content), registry.Locations[0].ID) || strings.Contains(string(registryInput.Content), registry.Locations[1].ID) || !strings.Contains(string(registryInput.Content), "registry_refs") {
		t.Fatalf("occurrence registry input = %#v, want contextual selector projection", registryInput)
	}
	contextArtifact := outputFileContent(t, output.OutputFiles, "evidence-context.json")
	evidence, err := evidencecontext.New().Decode(contextArtifact)
	if err != nil {
		t.Fatalf("Decode(evidence context) error = %v", err)
	}
	if actual := locationEvidenceUnitIDs(evidence); !reflect.DeepEqual(actual, []int{1, 2, 3, 4, 5}) {
		t.Fatalf("evidence context = %#v, want deduplicated registry and occurrence source-unit evidence", evidence)
	}
}

func locationEvidenceUnitIDs(document evidencecontext.Document) []int {
	ids := make([]int, len(document))
	for index, unit := range document {
		ids[index] = unit.ID
	}
	return ids
}

func TestLocationOccurrenceConsumerDoesNotRunAfterRejectedRegistry(t *testing.T) {
	registries := productionNPCRegistries(t)
	client := &locationHandoffLLMClient{rejectRegistry: true}
	_, err := runPreparedPipeline(t, registries, resolveLocationPipeline(t, registries), client, pipeline.RunInput{RawInput: readNPCFixture(t)})
	if err == nil {
		t.Fatal("Run() error = nil, want missing accepted location registry")
	}
	if client.requestCount(locationoccurrences.PromptID) != 0 {
		t.Fatalf("occurrence requests = %d, want none after rejected registry", client.requestCount(locationoccurrences.PromptID))
	}
}

func TestLocationOccurrenceConfigurationRequiresEarlierRegistry(t *testing.T) {
	registries := productionNPCRegistries(t)
	cfg, err := configFromYAML(strings.Replace(locationPipelineConfig, "step: identify-locations", "step: track-location-occurrences", 1))
	if err != nil {
		t.Fatal(err)
	}
	_, err = cfg.Resolve(config.ResolveInput{PipelineID: "dnd-locations-fixture", Catalog: moduleCatalog(registries)})
	if err == nil || !strings.Contains(err.Error(), "earlier step") {
		t.Fatalf("Resolve() error = %v, want an earlier generated registry requirement", err)
	}
}

func TestLocationOccurrenceCheckpointTracksGeneratedRegistry(t *testing.T) {
	registries := productionNPCRegistries(t)
	checkpoint := newGeneratedReferenceCheckpointLoader()
	_, err := runPreparedPipeline(t, registries, resolveLocationPipeline(t, registries), &locationHandoffLLMClient{}, pipeline.RunInput{
		RawInput:   readNPCFixture(t),
		Checkpoint: checkpoint,
	})
	if err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	for _, dependency := range checkpoint.extractDependencies("occurrences") {
		if dependency.Name == "generated-reference:location_registry:0" && strings.HasPrefix(dependency.Value, "sha256:") {
			return
		}
	}
	t.Fatalf("occurrence checkpoint dependencies = %#v, want generated location registry fingerprint", checkpoint.extractDependencies("occurrences"))
}

func resolveLocationPipeline(t *testing.T, registries pipeline.Registries) pipeline.ResolvedPipeline {
	t.Helper()
	cfg, err := configFromYAML(locationPipelineConfig)
	if err != nil {
		t.Fatal(err)
	}
	effective, err := cfg.Resolve(config.ResolveInput{PipelineID: "dnd-locations-fixture", Catalog: moduleCatalog(registries)})
	if err != nil {
		t.Fatalf("Resolve() error = %v", err)
	}
	resolved, warnings, err := pipeline.MaterializeReferences(effective.ResolvedPipeline, moduleCatalog(registries), pipeline.ReferenceMaterializationOptions{})
	if err != nil || len(warnings) != 0 {
		t.Fatalf("MaterializeReferences() error = %v warnings = %#v", err, warnings)
	}
	return resolved
}

func configFromYAML(content string) (config.Config, error) {
	fileConfig, err := config.ParseFileConfigYAML([]byte(content))
	if err != nil {
		return config.Config{}, err
	}
	cfg := config.Default()
	if err := cfg.ApplyFileConfig(fileConfig); err != nil {
		return config.Config{}, err
	}
	return cfg, nil
}

const locationPipelineConfig = `version: 4
cache:
  chunk_plans:
    mode: bypass
  checkpoints: {}
pipelines:
  dnd-locations-fixture:
    input: seriatim
    output:
      module: json
      options:
        evidence_context:
          enabled: true
          lanes: [locations, occurrences]
    steps:
      - id: identify-locations
        artifacts:
          locations:
            extract: dnd/location-registry
            merge: appendorder
            normalize: dnd/location-registry
      - id: track-location-occurrences
        references:
          location_registry:
            artifact:
              step: identify-locations
              lane: locations
        artifacts:
          occurrences:
            extract: dnd/location-occurrences
            merge: appendorder
            normalize: dnd/location-occurrences
`

type locationHandoffLLMClient struct {
	mu             sync.Mutex
	requests       []contracts.StructuredCompletionRequest
	rejectRegistry bool
}

func (client *locationHandoffLLMClient) CompleteStructured(ctx context.Context, request contracts.StructuredCompletionRequest, output any) (contracts.StructuredCompletionResponse, error) {
	if err := ctx.Err(); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	client.mu.Lock()
	client.requests = append(client.requests, cloneStructuredCompletionRequest(request))
	client.mu.Unlock()
	var payload any
	switch request.PromptID {
	case locations.PromptID:
		name := "Moon Gate"
		if client.rejectRegistry {
			name = ""
		}
		payload = map[string]any{"locations": []any{map[string]any{"name": name, "source_refs": []any{map[string]int{"start_unit_id": 1, "end_unit_id": 1}}}, map[string]any{"name": name, "source_refs": []any{map[string]int{"start_unit_id": 3, "end_unit_id": 3}}}}}
	case locationnormalize.PromptID:
		payload = map[string]any{"duplicate_groups": []any{}}
	case locationoccurrences.PromptID:
		var projection struct {
			Locations []struct {
				Name         string `json:"name"`
				RegistryRefs []struct {
					StartUnitID int `json:"start_unit_id"`
					EndUnitID   int `json:"end_unit_id"`
				} `json:"registry_refs"`
			} `json:"locations"`
		}
		if err := json.Unmarshal(request.Inputs["location_registry"].Content, &projection); err != nil {
			return contracts.StructuredCompletionResponse{}, fmt.Errorf("decode location registry: %w", err)
		}
		if len(projection.Locations) != 2 {
			return contracts.StructuredCompletionResponse{}, fmt.Errorf("generated location registry has %d locations, want 2", len(projection.Locations))
		}
		payload = map[string]any{"occurrences": []any{map[string]any{"name": projection.Locations[0].Name, "registry_refs": projection.Locations[0].RegistryRefs, "kind": "visited", "source_refs": []any{map[string]int{"start_unit_id": 1, "end_unit_id": 1}}}, map[string]any{"name": projection.Locations[1].Name, "registry_refs": projection.Locations[1].RegistryRefs, "kind": "mentioned", "source_refs": []any{map[string]int{"start_unit_id": 3, "end_unit_id": 3}}}}}
	default:
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("unexpected prompt %q", request.PromptID)
	}
	content, err := json.Marshal(payload)
	if err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	if err := json.Unmarshal(content, output); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	return contracts.StructuredCompletionResponse{Content: content, Provider: "test", Model: "location-handoff"}, nil
}

func (client *locationHandoffLLMClient) requestFor(t *testing.T, promptID string) contracts.StructuredCompletionRequest {
	t.Helper()
	client.mu.Lock()
	defer client.mu.Unlock()
	for _, request := range client.requests {
		if request.PromptID == promptID {
			return request
		}
	}
	t.Fatalf("requests = %#v, missing %q", client.requests, promptID)
	return contracts.StructuredCompletionRequest{}
}
func (client *locationHandoffLLMClient) requestCount(promptID string) int {
	client.mu.Lock()
	defer client.mu.Unlock()
	count := 0
	for _, request := range client.requests {
		if request.PromptID == promptID {
			count++
		}
	}
	return count
}

var _ contracts.StructuredLLMClient = (*locationHandoffLLMClient)(nil)