package cli

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"strings"
	"sync"
	"testing"

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

func TestRecomputeStepRecoversThroughFilesystemCheckpoints(t *testing.T) {
	tests := []struct {
		name             string
		invalidateOutput bool
		wantCode         int
	}{
		{name: "accepted producer is hydrated", wantCode: 0},
		{name: "invalid producer stops dependents", invalidateOutput: true, wantCode: 1},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			roots := newRecomputeTestRoots(t)
			harness := newRecomputeTestHarness()
			fresh := runRecomputeCommand(roots, harness.options(), false)
			if fresh.code != 0 {
				t.Fatalf("fresh run code=%d stderr=%q", fresh.code, fresh.stderr)
			}
			removeCheckpointLaneStage(t, roots.checkpoints, "extract", "first", "producer")
			removeCheckpointLaneStage(t, roots.checkpoints, "merge", "first", "producer")
			if tt.invalidateOutput {
				path := findCheckpointFile(t, roots.checkpoints, "normalize", "first", "producer", "output.json")
				if err := os.WriteFile(path, []byte("{"), 0o600); err != nil {
					t.Fatal(err)
				}
			}
			harness.resetCalls()

			resumed := runRecomputeCommand(roots, harness.options(), true)
			if resumed.code != tt.wantCode {
				t.Fatalf("resumed code=%d stdout=%q stderr=%q", resumed.code, resumed.stdout, resumed.stderr)
			}
			events := readLatestCheckpointEvents(t, roots.debug)
			if tt.invalidateOutput {
				if harness.callsFor("test/extract/middle") != 0 || harness.callsFor("test/extract/dependent") != 0 {
					t.Fatalf("dependent calls after invalid producer = %#v", harness.callsSnapshot())
				}
				if !strings.Contains(resumed.stderr, string(pipeline.CheckpointReasonDecodeFailed)) {
					t.Fatalf("stderr=%q, want stable checkpoint reason", resumed.stderr)
				}
				assertNormalizeDecisionSequence(t, events, []checkpointDecisionExpectation{{"first", "producer", pipeline.CheckpointDecisionExecuted, pipeline.CheckpointReasonDecodeFailed}})
				return
			}

			if got := harness.callsSnapshot(); !reflect.DeepEqual(got, map[string]int{"test/extract/dependent": 1, "test/extract/middle": 1}) {
				t.Fatalf("resumed extractor calls = %#v", got)
			}
			outputPath := filepath.Join(latestChildDir(t, roots.output), "result.json")
			data, err := os.ReadFile(outputPath)
			if err != nil {
				t.Fatal(err)
			}
			if string(data) != "[\"producer\",\"unrelated\",\"middle\",\"dependent\"]\n" {
				t.Fatalf("ordered output = %q", data)
			}
			assertNormalizeDecisionSequence(t, events, []checkpointDecisionExpectation{
				{"first", "producer", pipeline.CheckpointDecisionReused, pipeline.CheckpointReasonAcceptedArtifactReused},
				{"first", "unrelated", pipeline.CheckpointDecisionReused, pipeline.CheckpointReasonReused},
				{"second", "middle", pipeline.CheckpointDecisionForcedRecompute, pipeline.CheckpointReasonRecomputeStep},
				{"third", "dependent", pipeline.CheckpointDecisionForcedRecompute, pipeline.CheckpointReasonRecomputeStep},
			})
		})
	}
}

type checkpointDecisionExpectation struct {
	step, lane string
	action     pipeline.CheckpointDecisionCategory
	reason     pipeline.CheckpointReasonCode
}

func assertNormalizeDecisionSequence(t *testing.T, events []pipeline.CheckpointEvent, want []checkpointDecisionExpectation) {
	t.Helper()
	var got []checkpointDecisionExpectation
	for _, event := range events {
		if event.Stage == string(pipeline.StageNormalize) {
			got = append(got, checkpointDecisionExpectation{event.StepID, event.LaneID, event.Action, event.ReasonCode})
		}
	}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("normalize decisions = %#v, want %#v", got, want)
	}
}

type recomputeTestHarness struct {
	base  *stateTestHarness
	mu    sync.Mutex
	calls map[string]int
}

func newRecomputeTestHarness() *recomputeTestHarness {
	return &recomputeTestHarness{base: newStateTestHarness(), calls: make(map[string]int)}
}

func (h *recomputeTestHarness) options() Options {
	opts := h.base.options()
	for _, key := range []string{"test/extract/producer", "test/extract/unrelated", "test/extract/middle", "test/extract/dependent"} {
		moduleKey := key
		spec := pipeline.ModuleSpec{
			Key: moduleKey, Stage: pipeline.StageExtract, Requires: []string{"chunks"}, Provides: []string{"artifact"}, ArtifactKind: stateTestArtifactKind,
			ReferenceSlots: []contracts.ReferenceSlot{{Name: "upstream", AcceptedMediaTypes: []string{"application/json"}, AcceptedArtifactKinds: []contracts.ArtifactKind{stateTestArtifactKind}}},
		}
		if err := pipeline.RegisterExtractor(opts.Registries.Extractors, spec, func() (contracts.Extractor[stateTestArtifact], error) {
			return recomputeTestExtractor{key: moduleKey, harness: h}, nil
		}); err != nil {
			panic(err)
		}
	}
	if err := opts.Registries.Outputs.RegisterWithSpec(pipeline.ModuleSpec{Key: "test/recompute-output", Stage: pipeline.StageOutput, Requires: []string{"normalized"}, Provides: []string{"output"}}, func() (contracts.OutputEncoder, error) {
		return recomputeTestOutput{}, nil
	}); err != nil {
		panic(err)
	}
	opts.Catalog = catalogFromRegistries(opts.Registries)
	return opts
}

func (h *recomputeTestHarness) record(key string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.calls[key]++
}

func (h *recomputeTestHarness) resetCalls() {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.calls = make(map[string]int)
}

func (h *recomputeTestHarness) callsFor(key string) int {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.calls[key]
}

func (h *recomputeTestHarness) callsSnapshot() map[string]int {
	h.mu.Lock()
	defer h.mu.Unlock()
	result := make(map[string]int, len(h.calls))
	for key, value := range h.calls {
		result[key] = value
	}
	return result
}

type recomputeTestExtractor struct {
	key     string
	harness *recomputeTestHarness
}

func (e recomputeTestExtractor) Key() string { return e.key }
func (e recomputeTestExtractor) ReferenceSlots() []contracts.ReferenceSlot {
	return []contracts.ReferenceSlot{{Name: "upstream", AcceptedMediaTypes: []string{"application/json"}, AcceptedArtifactKinds: []contracts.ArtifactKind{stateTestArtifactKind}}}
}
func (e recomputeTestExtractor) Extract(context.Context, contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[stateTestArtifact], error) {
	e.harness.record(e.key)
	return contracts.TypedExtractionResult[stateTestArtifact]{Value: stateTestArtifact{Value: e.key}}, nil
}

type recomputeTestOutput struct{}

func (recomputeTestOutput) Key() string { return "test/recompute-output" }
func (recomputeTestOutput) Encode(_ context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
	lanes := make([]string, 0, len(req.NormalizeOutputs))
	for _, output := range req.NormalizeOutputs {
		lanes = append(lanes, output.LaneID)
	}
	data, err := json.Marshal(lanes)
	if err != nil {
		return contracts.OutputResult{}, err
	}
	return contracts.OutputResult{Files: []contracts.OutputFile{{Name: "result.json", Bytes: append(data, '\n')}}}, nil
}

func newRecomputeTestRoots(t *testing.T) stateTestRoots {
	t.Helper()
	roots := newStateTestRoots(t)
	config := fmt.Sprintf(`version: 3
output:
  directory: %q
cache:
  chunk_plans:
    directory: %q
    mode: bypass
  checkpoints:
    enabled: true
    directory: %q
debug:
  directory: %q
pipelines:
  sample:
    input: test/input
    chunk: test/chunk
    steps:
      - id: first
        artifacts:
          producer:
            extract: test/extract/producer
            merge: test/merge
            normalize: test/normalize
          unrelated:
            extract: test/extract/unrelated
            merge: test/merge
            normalize: test/normalize
      - id: second
        references:
          upstream:
            artifact:
              step: first
              lane: producer
        artifacts:
          middle:
            extract: test/extract/middle
            merge: test/merge
            normalize: test/normalize
      - id: third
        references:
          upstream:
            artifact:
              step: second
              lane: middle
        artifacts:
          dependent:
            extract: test/extract/dependent
            merge: test/merge
            normalize: test/normalize
    output: test/recompute-output
`, roots.output, roots.plans, roots.checkpoints, roots.debug)
	if err := os.WriteFile(roots.config, []byte(config), 0o600); err != nil {
		t.Fatal(err)
	}
	return roots
}

func runRecomputeCommand(roots stateTestRoots, opts Options, recompute bool) stateTestResult {
	args := []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass"}
	if recompute {
		args = append(args, "--resume", "--recompute-step", "second", "--debug")
	}
	var stdout, stderr bytes.Buffer
	return stateTestResult{code: RunWithOptions(args, &stdout, &stderr, opts), stdout: stdout.String(), stderr: stderr.String()}
}

func removeCheckpointLaneStage(t *testing.T, root, stage, step, lane string) {
	t.Helper()
	dir := filepath.Dir(findCheckpointFile(t, root, stage, step, lane, "manifest.json"))
	if err := os.RemoveAll(dir); err != nil {
		t.Fatal(err)
	}
}

func findCheckpointFile(t *testing.T, root, stage, step, lane, name string) string {
	t.Helper()
	want := filepath.Join(stage, step, lane, name)
	var matches []string
	err := filepath.WalkDir(root, func(path string, entry os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !entry.IsDir() && strings.HasSuffix(path, want) {
			matches = append(matches, path)
		}
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	if len(matches) != 1 {
		t.Fatalf("checkpoint files ending in %q = %v", want, matches)
	}
	return matches[0]
}

func latestChildDir(t *testing.T, root string) string {
	t.Helper()
	entries, err := os.ReadDir(root)
	if err != nil {
		t.Fatal(err)
	}
	var dirs []string
	for _, entry := range entries {
		if entry.IsDir() {
			dirs = append(dirs, filepath.Join(root, entry.Name()))
		}
	}
	if len(dirs) == 0 {
		t.Fatal("no child directory")
	}
	sort.Strings(dirs)
	return dirs[len(dirs)-1]
}

func readLatestCheckpointEvents(t *testing.T, root string) []pipeline.CheckpointEvent {
	t.Helper()
	var events []pipeline.CheckpointEvent
	data, err := os.ReadFile(filepath.Join(latestChildDir(t, root), "summary", "checkpoint-events.json"))
	if err != nil {
		t.Fatal(err)
	}
	if err := json.Unmarshal(data, &events); err != nil {
		t.Fatal(err)
	}
	return events
}