package api

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"hash"
	"io"
	"sync"
	"sync/atomic"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/stretchr/testify/require"
)

// webSocketMessageCopier is the deterministic per-message core of the P3
// synchronous pump. Production loops over gorilla NextReader/NextWriter and
// delegates one message at a time to this method. openTarget defers the
// writer opening until the first non-empty Read clears the cumulative
// limit check (P4 lazy-open directive #2); the standard io types keep
// this structural red test buildable before the method exists.
type webSocketMessageCopier interface {
	copyWebSocketMessage(
		source io.Reader,
		openTarget func() (io.WriteCloser, error),
		maxMessageBytes int64,
	) error
}

type webSocketCloseSanitizer interface {
	sanitizeWebSocketClose(code int, reason string) (int, string)
}

func requireWebSocketMessageCopier(t *testing.T) webSocketMessageCopier {
	t.Helper()
	copier, ok := any(&desktopGateway{}).(webSocketMessageCopier)
	if !ok {
		t.Fatal("desktopGateway does not implement bounded WebSocket message copy")
	}
	return copier
}

type recordingMessageWriter struct {
	bytes      bytes.Buffer
	closeCalls atomic.Int32
}

func (w *recordingMessageWriter) Write(p []byte) (int, error) {
	return w.bytes.Write(p)
}

func (w *recordingMessageWriter) Close() error {
	w.closeCalls.Add(1)
	return nil
}

type dataThenErrorReader struct {
	data []byte
	err  error
	done bool
}

func (r *dataThenErrorReader) Read(p []byte) (int, error) {
	if r.done {
		return 0, io.EOF
	}
	r.done = true
	return copy(p, r.data), r.err
}

// singleUseOpener returns an opener that hands out target exactly once.
// The unit tests do not exercise multi-open paths; production calls the
// opener at most once per message (either when the first non-empty Read
// clears the limit, or — for empty messages — at EOF before Close).
func singleUseOpener(target io.WriteCloser) func() (io.WriteCloser, error) {
	return func() (io.WriteCloser, error) { return target, nil }
}

func TestGatewayWebSocketMessagePartialReadErrorDoesNotFinishTruncatedMessage(t *testing.T) {
	copier := requireWebSocketMessageCopier(t)
	source := &dataThenErrorReader{
		data: []byte("partial"),
		err:  io.ErrUnexpectedEOF,
	}
	target := &recordingMessageWriter{}

	err := copier.copyWebSocketMessage(source, singleUseOpener(target), 8<<20)
	if !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Fatalf("copy error = %v, want UnexpectedEOF", err)
	}
	if got := target.bytes.String(); got != "partial" {
		t.Fatalf("partial forwarded bytes = %q", got)
	}
	if got := target.closeCalls.Load(); got != 0 {
		t.Fatalf("writer Close calls = %d, want 0 so no FIN commits truncated message", got)
	}
}

func TestGatewayWebSocketMessageOversizeDoesNotFinishMessage(t *testing.T) {
	const limit = int64(1024)
	copier := requireWebSocketMessageCopier(t)
	source := bytes.NewReader(bytes.Repeat([]byte{'x'}, int(limit+1)))
	target := &recordingMessageWriter{}
	var opened atomic.Bool
	opener := func() (io.WriteCloser, error) {
		opened.Store(true)
		return target, nil
	}

	err := copier.copyWebSocketMessage(source, opener, limit)
	if err == nil {
		t.Fatal("oversized message returned nil error")
	}
	if opened.Load() {
		t.Fatal("lazy-open contract violated: target writer was opened on oversize")
	}
	if got := target.closeCalls.Load(); got != 0 {
		t.Fatalf("writer Close calls = %d, want 0 on oversize", got)
	}
	if got := int64(target.bytes.Len()); got > limit {
		t.Fatalf("forwarded bytes = %d, exceeds hard limit %d", got, limit)
	}
}

func TestGatewayWebSocketMessageNormalEOFClosesExactlyOnce(t *testing.T) {
	copier := requireWebSocketMessageCopier(t)
	payload := []byte("complete")
	target := &recordingMessageWriter{}

	err := copier.copyWebSocketMessage(bytes.NewReader(payload), singleUseOpener(target), 8<<20)
	if err != nil {
		t.Fatalf("copy: %v", err)
	}
	if !bytes.Equal(target.bytes.Bytes(), payload) {
		t.Fatalf("payload = %q, want %q", target.bytes.Bytes(), payload)
	}
	if got := target.closeCalls.Load(); got != 1 {
		t.Fatalf("writer Close calls = %d, want exactly 1", got)
	}
}

func TestGatewayWebSocketCloseReasonSanitizesInvalidOrOversizedData(t *testing.T) {
	sanitizer, ok := any(&desktopGateway{}).(webSocketCloseSanitizer)
	if !ok {
		t.Fatal("desktopGateway does not implement WebSocket close sanitization")
	}
	cases := []struct {
		name       string
		code       int
		reason     string
		wantCode   int
		wantReason string
	}{
		{name: "valid", code: 1000, reason: "normal", wantCode: 1000, wantReason: "normal"},
		{name: "invalid-utf8", code: 1000, reason: string([]byte{0xff}), wantCode: 1000},
		{name: "too-long", code: 1000, reason: string(bytes.Repeat([]byte{'x'}, 124)), wantCode: 1000},
		{name: "reserved-code", code: 1005, reason: "hidden", wantCode: 1000},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			code, reason := sanitizer.sanitizeWebSocketClose(tc.code, tc.reason)
			if code != tc.wantCode || reason != tc.wantReason {
				t.Fatalf("sanitize = (%d,%q), want (%d,%q)", code, reason, tc.wantCode, tc.wantReason)
			}
			if !utf8.ValidString(reason) || len([]byte(reason)) > 123 {
				t.Fatalf("sanitized reason is invalid: %q", reason)
			}
		})
	}
}

// gatingTargetWriter blocks its first Write on a release channel and
// counts every Write/Close call. It also accumulates a sha256 hash
// over the bytes it accepts, so the test can prove end-to-end payload
// integrity after release — not just byte counts. It is the
// deterministic twin of the 1MiB-socket-plus-deadline inference
// pattern: the test observes exact Read/Write counts at the moment
// Write parks and compares the real hash after release, replacing
// timing and buffer-size assumptions with channel barriers.
type gatingTargetWriter struct {
	writeCalls   atomic.Int32
	closeCalls   atomic.Int32
	writeEntered chan struct{}
	releaseWrite chan struct{}
	writeOnce    sync.Once
	hasher       hash.Hash
	hashMu       sync.Mutex
}

func (w *gatingTargetWriter) Write(p []byte) (int, error) {
	w.writeCalls.Add(1)
	w.writeOnce.Do(func() { close(w.writeEntered) })
	<-w.releaseWrite
	if len(p) > 0 {
		w.hashMu.Lock()
		_, _ = w.hasher.Write(p)
		w.hashMu.Unlock()
	}
	return len(p), nil
}

// sumHex returns the hex-encoded sha256 of every byte the writer
// accepted after its gate released. The test computes the matching
// sha256 over the source chunks and the two must match exactly.
func (w *gatingTargetWriter) sumHex() string {
	w.hashMu.Lock()
	defer w.hashMu.Unlock()
	return hex.EncodeToString(w.hasher.Sum(nil))
}

func (w *gatingTargetWriter) Close() error {
	w.closeCalls.Add(1)
	return nil
}

// oneShotSource hands out exactly one non-empty Read and then EOF. The
// read counter lets the gated Write test prove the pump has not read
// past the first chunk while the target's Write is parked.
type oneShotSource struct {
	readCalls atomic.Int32
}

func (s *oneShotSource) Read(p []byte) (int, error) {
	n := s.readCalls.Add(1)
	if n == 1 {
		return copy(p, []byte("hello")), nil
	}
	return 0, io.EOF
}

// TestGatewayWebSocketMessageGatedTargetWriteBlocksSourceRead is the
// deterministic within-message backpressure test. While the target's
// first Write is parked on its release channel, the source MUST NOT be
// read a second time — otherwise a slow target would let the pump
// buffer unbounded source data. Releasing the gate completes the
// message and the pump closes the target exactly once.
//
// This replaces the v2 inference pattern (1MiB socket + 100ms
// deadline) with an explicit writer gate.
func TestGatewayWebSocketMessageGatedTargetWriteBlocksSourceRead(t *testing.T) {
	copier := requireWebSocketMessageCopier(t)
	source := &oneShotSource{}
	target := newGatingTargetWriter()
	done := make(chan error, 1)
	go func() {
		done <- copier.copyWebSocketMessage(source, singleUseOpener(target), 8<<20)
	}()

	select {
	case <-target.writeEntered:
	case <-time.After(time.Second):
		t.Fatal("target.Write was never called")
	}
	require.Equal(t, int32(1), source.readCalls.Load(),
		"source must not be read again while target.Write is parked")
	require.Equal(t, int32(1), target.writeCalls.Load(),
		"target.Write must be called exactly once at the gate")
	require.Equal(t, int32(0), target.closeCalls.Load(),
		"target must not be closed before the message completes")

	close(target.releaseWrite)
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("copyWebSocketMessage: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("copy did not complete after target.Write released")
	}
	require.Equal(t, int32(1), target.closeCalls.Load(),
		"target must be closed exactly once after EOF")
}

// failingTargetWriter fails every Write. The pump's no-FIN guarantee
// relies on Close NOT being called when Write errors mid-message; the
// WS frame FIN bit is set by Close, so a Close-on-error would commit
// a truncated frame to the wire.
type failingTargetWriter struct {
	writeCalls atomic.Int32
	closeCalls atomic.Int32
}

func (w *failingTargetWriter) Write(p []byte) (int, error) {
	w.writeCalls.Add(1)
	return 0, errors.New("target write failed")
}

func (w *failingTargetWriter) Close() error {
	w.closeCalls.Add(1)
	return nil
}

func newGatingTargetWriter() *gatingTargetWriter {
	return &gatingTargetWriter{
		writeEntered: make(chan struct{}),
		releaseWrite: make(chan struct{}),
		hasher:       sha256.New(),
	}
}

// hashChunks returns the hex sha256 of the concatenated chunks, matching
// what gatingTargetWriter.sumHex would observe after accepting the
// same byte sequence.
func hashChunks(chunks [][]byte) string {
	h := sha256.New()
	for _, chunk := range chunks {
		_, _ = h.Write(chunk)
	}
	return hex.EncodeToString(h.Sum(nil))
}

// TestGatewayWebSocketMessageTargetWriteErrorAbortsNoFIN locks the
// no-FIN-on-error invariant: when the target's Write returns an error
// mid-message, the pump MUST propagate the error WITHOUT calling
// Close. WS frame FIN is set by gorilla when the writer is closed, so
// a Close-on-error would deliver a truncated frame to the peer.
func TestGatewayWebSocketMessageTargetWriteErrorAbortsNoFIN(t *testing.T) {
	copier := requireWebSocketMessageCopier(t)
	source := bytes.NewReader([]byte("hello"))
	target := &failingTargetWriter{}

	err := copier.copyWebSocketMessage(source, singleUseOpener(target), 8<<20)
	if err == nil {
		t.Fatal("copyWebSocketMessage returned nil error on target Write failure")
	}
	require.Equal(t, int32(1), target.writeCalls.Load(),
		"target.Write must be attempted at least once")
	require.Equal(t, int32(0), target.closeCalls.Load(),
		"target.Close must NOT be called on Write error — no FIN commits a truncated frame")
}

// gatedHashSource emits deterministic chunks so the released-write
// variant can prove payload integrity end-to-end: the bytes the target
// receives must hash to exactly what the source produced, regardless
// of how long the Write was parked.
type gatedHashSource struct {
	chunks    [][]byte
	readCalls atomic.Int32
}

func (s *gatedHashSource) Read(p []byte) (int, error) {
	n := int(s.readCalls.Add(1))
	if n > len(s.chunks) {
		return 0, io.EOF
	}
	return copy(p, s.chunks[n-1]), nil
}

// TestGatewayWebSocketMessageGatedTargetWritePreservesPayloadHash
// complements the gated-Write backpressure test: once the gate is
// released, the full payload must arrive intact and the target is
// closed exactly once. The writer accumulates a sha256 over every
// byte accepted after the gate; the source's chunks hash must match
// exactly. This replaces the v3 read-count-only assertion with real
// end-to-end payload integrity.
func TestGatewayWebSocketMessageGatedTargetWritePreservesPayloadHash(t *testing.T) {
	copier := requireWebSocketMessageCopier(t)
	chunks := [][]byte{
		bytes.Repeat([]byte("alpha-"), 32),
		bytes.Repeat([]byte("beta-"), 16),
	}
	source := &gatedHashSource{chunks: chunks}
	target := newGatingTargetWriter()
	done := make(chan error, 1)
	go func() {
		done <- copier.copyWebSocketMessage(source, singleUseOpener(target), 8<<20)
	}()

	<-target.writeEntered
	close(target.releaseWrite)

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("copyWebSocketMessage: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("copy did not complete after gate release")
	}

	require.Equal(t, int32(1), target.closeCalls.Load(),
		"target must be closed exactly once after EOF")
	require.Equal(t, int32(len(chunks)+1), source.readCalls.Load(),
		"source must be read once per chunk plus a final EOF")
	require.Equal(t, hashChunks(chunks), target.sumHex(),
		"target must accept the exact source payload after gate release; "+
			"writer must not discard or mutate bytes")
}