package validate

import (
	"context"
	"errors"
	"io/fs"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"testing/fstest"
	"time"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"github.com/santhosh-tekuri/jsonschema/v6"
)

func TestValidationCancellationBeforeWorkDoesNotOpenSchemaSource(t *testing.T) {
	source := &countingSchemaFS{FS: fstest.MapFS{
		"schema.json": {Data: []byte(`{"type":"object"}`)},
	}}
	validator := NewFSValidator(source, ".").(ValidationPreparer)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
		ValidationMode: domain.ValidationJSONSchema,
		SchemaPath:     "schema.json",
	})
	if plan != nil || !errors.Is(err, context.Canceled) {
		t.Fatalf("PrepareValidation() = (%v, %v), want nil plan and context cancellation", plan, err)
	}
	if opens := source.opens.Load(); opens != 0 {
		t.Fatalf("schema source opens = %d, want 0", opens)
	}
}

func TestValidationCancellationBetweenReferencedSchemaReadChunks(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	source := &controlledSchemaFS{
		FS: fstest.MapFS{
			"root.json":  {Data: []byte(`{"$ref":"child.json"}`)},
			"child.json": {Data: []byte(`{"type":"string"}` + strings.Repeat(" ", schemaReadChunkSize*2))},
		},
		target: "child.json",
		cancel: cancel,
	}
	validator := NewFSValidator(source, ".").(ValidationPreparer)

	plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
		ValidationMode: domain.ValidationJSONSchema,
		SchemaPath:     "root.json",
	})
	if plan != nil || !errors.Is(err, context.Canceled) {
		t.Fatalf("PrepareValidation() = (%v, %v), want nil plan and context cancellation", plan, err)
	}
	if reads := source.reads.Load(); reads != 1 {
		t.Fatalf("controlled child reads = %d, want 1", reads)
	}
	if closes := source.closes.Load(); closes != 1 {
		t.Fatalf("controlled child closes = %d, want 1", closes)
	}
}

func TestDecodeCancellationAfterSynchronousCallWins(t *testing.T) {
	ctx := newCheckpointContext(2)
	value, err := decodeJSONValue(ctx, []byte(`{"value":1}`))
	if value != nil || !errors.Is(err, context.Canceled) {
		t.Fatalf("decodeJSONValue() = (%v, %v), want nil value and context cancellation", value, err)
	}
}

func TestCompileCancellationAfterSynchronousCallWins(t *testing.T) {
	for _, dependencyErr := range []error{nil, errors.New("compile failed")} {
		ctx, cancel := context.WithCancel(context.Background())
		compiled := &jsonschema.Schema{}
		schema, err := compileJSONSchema(ctx, "promptkit-schema:/root.json", func(string) (*jsonschema.Schema, error) {
			cancel()
			return compiled, dependencyErr
		})
		if schema != nil || !errors.Is(err, context.Canceled) {
			t.Fatalf("compileJSONSchema() = (%v, %v), want authoritative context cancellation", schema, err)
		}
		if dependencyErr != nil && errors.Is(err, dependencyErr) {
			t.Fatalf("compileJSONSchema() error = %v, dependency error should not win", err)
		}
	}
}

func TestExecutionCancellationAfterSynchronousCallWins(t *testing.T) {
	for _, dependencyErr := range []error{nil, errors.New("schema mismatch")} {
		ctx, cancel := context.WithCancel(context.Background())
		executor := &controlledSchemaExecutor{cancel: cancel, err: dependencyErr}
		validationErrors, err := executeJSONSchema(ctx, executor, map[string]any{"ok": true})
		if validationErrors != nil || !errors.Is(err, context.Canceled) {
			t.Fatalf("executeJSONSchema() = (%v, %v), want no result and context cancellation", validationErrors, err)
		}
		if calls := executor.calls.Load(); calls != 1 {
			t.Fatalf("schema execution calls = %d, want 1", calls)
		}
	}
}

func TestCancellationDoesNotDetachBlockedSchemaRead(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	started := make(chan struct{})
	release := make(chan struct{})
	source := &controlledSchemaFS{
		FS: fstest.MapFS{
			"schema.json": {Data: []byte(`{"type":"object"}`)},
		},
		target:  "schema.json",
		started: started,
		release: release,
	}
	validator := NewFSValidator(source, ".").(ValidationPreparer)
	type outcome struct {
		plan PreparedValidation
		err  error
	}
	result := make(chan outcome, 1)
	go func() {
		plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
			ValidationMode: domain.ValidationJSONSchema,
			SchemaPath:     "schema.json",
		})
		result <- outcome{plan: plan, err: err}
	}()

	waitForSignal(t, started, "schema read to start")
	cancel()
	select {
	case got := <-result:
		t.Fatalf("blocked dependency returned before release: (%v, %v)", got.plan, got.err)
	default:
	}
	close(release)
	got := waitForOutcome(t, result)
	if got.plan != nil || !errors.Is(got.err, context.Canceled) {
		t.Fatalf("PrepareValidation() after release = (%v, %v), want nil plan and context cancellation", got.plan, got.err)
	}
	if reads := source.reads.Load(); reads != 1 {
		t.Fatalf("blocked reads = %d, want 1 completed read", reads)
	}
}

func TestCancellationDoesNotDetachBlockedSchemaExecution(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	started := make(chan struct{})
	release := make(chan struct{})
	executor := &controlledSchemaExecutor{started: started, release: release}
	result := make(chan error, 1)
	go func() {
		_, err := executeJSONSchema(ctx, executor, map[string]any{"ok": true})
		result <- err
	}()

	waitForSignal(t, started, "schema execution to start")
	cancel()
	select {
	case err := <-result:
		t.Fatalf("blocked dependency returned before release: %v", err)
	default:
	}
	close(release)
	select {
	case err := <-result:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("executeJSONSchema() after release = %v, want context cancellation", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for released schema execution")
	}
	if calls := executor.calls.Load(); calls != 1 {
		t.Fatalf("schema execution calls = %d, want 1 completed call", calls)
	}
}

type countingSchemaFS struct {
	fs.FS
	opens atomic.Int32
}

func (f *countingSchemaFS) Open(name string) (fs.File, error) {
	f.opens.Add(1)
	return f.FS.Open(name)
}

type controlledSchemaFS struct {
	fs.FS
	target  string
	cancel  context.CancelFunc
	started chan struct{}
	release chan struct{}
	once    sync.Once
	reads   atomic.Int32
	closes  atomic.Int32
}

func (f *controlledSchemaFS) Open(name string) (fs.File, error) {
	file, err := f.FS.Open(name)
	if err != nil || name != f.target {
		return file, err
	}
	return &controlledSchemaFile{File: file, owner: f}, nil
}

type controlledSchemaFile struct {
	fs.File
	owner *controlledSchemaFS
}

func (f *controlledSchemaFile) Read(buffer []byte) (int, error) {
	f.owner.once.Do(func() {
		if f.owner.started != nil {
			close(f.owner.started)
		}
		if f.owner.release != nil {
			<-f.owner.release
		}
	})
	n, err := f.File.Read(buffer)
	f.owner.reads.Add(1)
	if f.owner.cancel != nil {
		f.owner.cancel()
	}
	return n, err
}

func (f *controlledSchemaFile) Close() error {
	f.owner.closes.Add(1)
	return f.File.Close()
}

type controlledSchemaExecutor struct {
	cancel  context.CancelFunc
	err     error
	started chan struct{}
	release chan struct{}
	calls   atomic.Int32
}

func (e *controlledSchemaExecutor) Validate(any) error {
	e.calls.Add(1)
	if e.started != nil {
		close(e.started)
	}
	if e.release != nil {
		<-e.release
	}
	if e.cancel != nil {
		e.cancel()
	}
	return e.err
}

type checkpointContext struct {
	context.Context
	mu        sync.Mutex
	remaining int
	canceled  bool
	done      chan struct{}
}

func newCheckpointContext(checksUntilCancel int) *checkpointContext {
	return &checkpointContext{
		Context:   context.Background(),
		remaining: checksUntilCancel,
		done:      make(chan struct{}),
	}
}

func (c *checkpointContext) Done() <-chan struct{} {
	return c.done
}

func (c *checkpointContext) Err() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.canceled {
		return context.Canceled
	}
	c.remaining--
	if c.remaining == 0 {
		c.canceled = true
		close(c.done)
		return context.Canceled
	}
	return nil
}

func waitForSignal(t *testing.T, signal <-chan struct{}, description string) {
	t.Helper()
	select {
	case <-signal:
	case <-time.After(2 * time.Second):
		t.Fatalf("timed out waiting for %s", description)
	}
}

func waitForOutcome[T any](t *testing.T, result <-chan T) T {
	t.Helper()
	select {
	case value := <-result:
		return value
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for released dependency")
		var zero T
		return zero
	}
}