package rpc

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/ws"
)

type memoryTransport struct {
	ctx   context.Context
	other *Peer
}

func (m *memoryTransport) Context() context.Context { return m.ctx }
func (m *memoryTransport) Send(ctx context.Context, msg ws.Message) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	return m.other.receive(msg)
}
func (m *memoryTransport) TrySend(msg ws.Message) error { return m.other.receive(msg) }
func (m *memoryTransport) Close(int, string) error      { return nil }
func (m *memoryTransport) Abort()                       { m.other.Stop() }

func pair(t testing.TB, cfg Config, setup func(*Peer, *Peer)) (*Peer, *Peer) {
	t.Helper()
	a, err := New(cfg)
	if err != nil {
		t.Fatal(err)
	}
	b, err := New(cfg)
	if err != nil {
		t.Fatal(err)
	}
	if setup != nil {
		setup(a, b)
	}
	ctx, cancel := context.WithCancel(context.Background())
	if err = a.bind(&memoryTransport{ctx, b}); err != nil {
		t.Fatal(err)
	}
	if err = b.bind(&memoryTransport{ctx, a}); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		cancel()
		a.Stop()
		b.Stop()
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		if err := a.Wait(ctx); err != nil {
			t.Error(err)
		}
		if err := b.Wait(ctx); err != nil {
			t.Error(err)
		}
	})
	return a, b
}
func handle(t testing.TB, p *Peer, name string, h Handler) {
	t.Helper()
	if err := p.Handle(name, h); err != nil {
		t.Fatal(err)
	}
}

func TestBidirectionalNestedCalls(t *testing.T) {
	a, _ := pair(t, Config{}, func(a, b *Peer) {
		handle(t, a, "double", func(_ context.Context, v json.RawMessage) (any, error) {
			var n int
			_ = json.Unmarshal(v, &n)
			return n * 2, nil
		})
		handle(t, b, "nested", func(ctx context.Context, v json.RawMessage) (any, error) {
			var n int
			_ = json.Unmarshal(v, &n)
			var doubled int
			err := b.Call(ctx, "double", n, &doubled)
			return doubled + 1, err
		})
	})
	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func(n int) {
			defer wg.Done()
			var result int
			if err := a.Call(context.Background(), "nested", n, &result); err != nil || result != n*2+1 {
				t.Errorf("result %d, error %v", result, err)
			}
		}(i)
	}
	wg.Wait()
}

func TestCancellationBusyAndCompletionBarrier(t *testing.T) {
	entered := make(chan struct{})
	exit := make(chan struct{})
	release := make(chan struct{})
	a, b := pair(t, Config{MaxConcurrent: 1}, func(a, b *Peer) {
		handle(t, b, "block", func(ctx context.Context, _ json.RawMessage) (any, error) {
			close(entered)
			<-ctx.Done()
			close(exit)
			<-release
			return nil, ctx.Err()
		})
	})
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	result := make(chan error, 1)
	go func() { result <- a.Call(ctx, "block", nil, nil) }()
	<-entered
	if err := a.Call(context.Background(), "block", nil, nil); err == nil {
		t.Fatal("capacity not enforced")
	} else {
		var pub *Error
		if !errors.As(err, &pub) || pub.Code != "BUSY" {
			t.Fatal(err)
		}
	}
	if err := <-result; !errors.Is(err, context.DeadlineExceeded) {
		t.Fatal(err)
	}
	<-exit
	b.Stop()
	wait, c := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer c()
	if err := b.Wait(wait); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatal("handler not joined", err)
	}
	close(release)
	if err := b.Wait(context.Background()); err != nil {
		t.Fatal(err)
	}
}

func TestErrorsDoNotExposeImplementation(t *testing.T) {
	a, _ := pair(t, Config{}, func(a, b *Peer) {
		handle(t, b, "private", func(context.Context, json.RawMessage) (any, error) { return nil, errors.New("password=secret") })
		handle(t, b, "panic", func(context.Context, json.RawMessage) (any, error) { panic("token=private") })
		handle(t, b, "public", func(context.Context, json.RawMessage) (any, error) {
			return nil, &Error{"DENIED", "public explanation"}
		})
	})
	for _, name := range []string{"private", "panic", "missing", "public"} {
		err := a.Call(context.Background(), name, nil, nil)
		var pub *Error
		if !errors.As(err, &pub) {
			t.Fatal(err)
		}
		if strings.Contains(pub.Message, "secret") || strings.Contains(pub.Message, "private") {
			t.Fatal("leaked error", pub.Message)
		}
		if name == "public" && pub.Message != "public explanation" {
			t.Fatal(pub)
		}
	}
}

func TestReplyLateDuplicateAndProtocolLimits(t *testing.T) {
	p, _ := New(Config{})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	q, _ := New(Config{})
	_ = p.bind(&memoryTransport{ctx, q})
	defer p.Stop()
	// Late and duplicated replies must not block the reader or create allocations.
	for i := 0; i < 10; i++ {
		if err := p.receive(ws.Message{Type: ws.Text, Data: []byte(`{"v":1,"type":"reply","id":"old","result":null}`)}); err != nil {
			t.Fatal(err)
		}
	}
	invalid := []string{
		`{"v":2,"type":"event","method":"x","params":null}`,
		`{"v":1,"type":"call","id":"x","id":"y","method":"x","params":null}`,
		`{"v":1,"type":"event","method":"x","params":{"role":1,"role":2}}`,
		`{"v":1,"type":"reply","id":"a","result":null,"error":{"code":"BAD","message":"x"}}`,
		`{"v":1,"type":"cancel","id":"a","params":null}`,
		`{"v":1,"type":"call","id":"x","method":"x","params":null,"extra":true}`,
		`{"v":1,"type":"event","method":"x","params":` + strings.Repeat("[", 34) + "0" + strings.Repeat("]", 34) + `}`,
	}
	for _, data := range invalid {
		if _, err := decode([]byte(data), 1<<20); !errors.Is(err, ErrProtocol) {
			t.Fatalf("accepted %s", data)
		}
	}
	if _, err := decode([]byte(`{"v":1,"type":"event","method":"x","params":null}`), 4); err == nil {
		t.Fatal("size ignored")
	}
	if err := p.receive(ws.Message{Type: ws.Binary, Data: []byte("x")}); !errors.Is(err, ErrProtocol) {
		t.Fatal(err)
	}
}

func TestPendingLimitAndDisconnect(t *testing.T) {
	entered := make(chan struct{})
	a, b := pair(t, Config{MaxPending: 1}, func(a, b *Peer) {
		handle(t, b, "wait", func(ctx context.Context, _ json.RawMessage) (any, error) {
			close(entered)
			<-ctx.Done()
			return nil, ctx.Err()
		})
	})
	done := make(chan error, 1)
	go func() { done <- a.Call(context.Background(), "wait", nil, nil) }()
	<-entered
	if err := a.Call(context.Background(), "wait", nil, nil); !errors.Is(err, ErrBusy) {
		t.Fatal(err)
	}
	a.Stop()
	b.Stop()
	if err := <-done; !errors.Is(err, ErrClosed) {
		t.Fatal(err)
	}
	if err := a.bind(&memoryTransport{context.Background(), b}); !errors.Is(err, ErrClosed) {
		t.Fatal("peer rebound", err)
	}
}

func TestSubscriptionsRestoreAndRemoval(t *testing.T) {
	var mu sync.Mutex
	var keys []string
	setup := func(a, b *Peer) {
		handle(t, b, "rpc.subscribe", func(_ context.Context, raw json.RawMessage) (any, error) {
			var sub Subscription
			if err := json.Unmarshal(raw, &sub); err != nil {
				return nil, err
			}
			mu.Lock()
			keys = append(keys, sub.Key)
			mu.Unlock()
			return nil, nil
		})
		handle(t, b, "rpc.unsubscribe", func(context.Context, json.RawMessage) (any, error) { return nil, nil })
	}
	a, _ := pair(t, Config{}, setup)
	s, err := NewSubscriptions(SubscriptionConfig{MaxSubscriptions: 2})
	if err != nil {
		t.Fatal(err)
	}
	_ = s.Subscribe(context.Background(), nil, "b", map[string]int{"n": 1})
	_ = s.Subscribe(context.Background(), nil, "a", nil)
	if err = s.Subscribe(context.Background(), nil, "c", nil); !errors.Is(err, ErrBusy) {
		t.Fatal(err)
	}
	if n, err := s.Restore(context.Background(), a); n != 2 || err != nil {
		t.Fatal(n, err)
	}
	if err = s.Unsubscribe(context.Background(), a, "a"); err != nil {
		t.Fatal(err)
	}
	next, _ := pair(t, Config{}, setup)
	if n, err := s.Restore(context.Background(), next); n != 1 || err != nil {
		t.Fatal(n, err)
	}
	mu.Lock()
	defer mu.Unlock()
	if strings.Join(keys, ",") != "a,b,b" {
		t.Fatal(keys)
	}
}

func TestRealTCPRPCAndShutdown(t *testing.T) {
	opened := make(chan *Peer, 1)
	server, err := ws.NewServer(ws.ServerConfig{Subprotocols: []string{Subprotocol}, RequireSubprotocol: true, Prepare: func(context.Context, *http.Request) (ws.Handlers, error) {
		p, _ := New(Config{})
		handle(t, p, "echo", func(_ context.Context, raw json.RawMessage) (any, error) { return raw, nil })
		opened <- p
		return p.Handlers(), nil
	}})
	if err != nil {
		t.Fatal(err)
	}
	httpServer := httptest.NewServer(server.Handler(ws.Handlers{}))
	defer httpServer.Close()
	client, err := ws.NewClient(ws.ClientConfig{Subprotocols: []string{Subprotocol}, RequireSubprotocol: true, ProxyMode: "direct"})
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		_ = client.Shutdown(ctx)
		_ = server.Shutdown(ctx)
	}()
	p, _ := New(Config{})
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	session, err := client.Dial(ctx, "ws"+strings.TrimPrefix(httpServer.URL, "http"), p.Handlers())
	if err != nil {
		t.Fatal(err)
	}
	if err = p.WaitReady(ctx); err != nil {
		t.Fatal(err)
	}
	remote := <-opened
	var result map[string]int
	if err = p.Call(ctx, "echo", map[string]int{"n": 42}, &result); err != nil || result["n"] != 42 {
		t.Fatal(result, err)
	}
	if err = p.Shutdown(ctx); err != nil {
		t.Fatal(err)
	}
	if err = session.Wait(ctx); err != nil {
		t.Fatal(err)
	}
	if err = remote.Wait(ctx); err != nil {
		t.Fatal(err)
	}
}

func BenchmarkPeerCall(b *testing.B) {
	a, _ := pair(b, Config{}, func(a, p *Peer) {
		handle(b, p, "echo", func(_ context.Context, raw json.RawMessage) (any, error) { return raw, nil })
	})
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		var out int
		if err := a.Call(context.Background(), "echo", 42, &out); err != nil {
			b.Fatal(err)
		}
	}
}
func BenchmarkProtocolDecode(b *testing.B) {
	data := []byte(`{"v":1,"type":"call","id":"abc","method":"status.get","params":{"detail":true},"timeoutMs":1000}`)
	b.ReportAllocs()
	b.SetBytes(int64(len(data)))
	for i := 0; i < b.N; i++ {
		if _, err := decode(data, 1<<20); err != nil {
			b.Fatal(err)
		}
	}
}
func FuzzProtocolDecode(f *testing.F) {
	for _, s := range []string{`{"v":1,"type":"event","method":"x","params":null}`, `{"v":1,"type":"reply","id":"x","result":42}`, `null`, `{"v":1,"v":2}`} {
		f.Add([]byte(s))
	}
	f.Fuzz(func(t *testing.T, data []byte) { _, _ = decode(data, 4096) })
}

type panickingResult struct{}

func (panickingResult) MarshalJSON() ([]byte, error) { panic("secret encoder state") }
func TestResultEncodingPanicAndOversize(t *testing.T) {
	a, _ := pair(t, Config{MaxMessageBytes: 256}, func(a, b *Peer) {
		handle(t, b, "panicResult", func(context.Context, json.RawMessage) (any, error) { return panickingResult{}, nil })
		handle(t, b, "large", func(context.Context, json.RawMessage) (any, error) { return strings.Repeat("x", 1024), nil })
	})
	for _, method := range []string{"panicResult", "large"} {
		err := a.Call(context.Background(), method, nil, nil)
		var e *Error
		if !errors.As(err, &e) {
			t.Fatal(err)
		}
		if method == "panicResult" && e.Code != "INTERNAL" {
			t.Fatal(e)
		}
		if method == "large" && e.Code != "RESULT_TOO_LARGE" {
			t.Fatal(e)
		}
	}
}
func TestEventErrorObservation(t *testing.T) {
	observed := make(chan error, 1)
	a, _ := pair(t, Config{OnEventError: func(_ context.Context, err error) { observed <- err; panic("observer panic") }}, func(a, b *Peer) {
		handle(t, b, "fail", func(context.Context, json.RawMessage) (any, error) { return nil, errors.New("local failure") })
	})
	if err := a.Notify(context.Background(), "fail", nil); err != nil {
		t.Fatal(err)
	}
	select {
	case err := <-observed:
		if err.Error() != "local failure" {
			t.Fatal(err)
		}
	case <-time.After(time.Second):
		t.Fatal("event error not observed")
	}
}
func TestDuplicateIncomingIDRejected(t *testing.T) {
	entered := make(chan struct{})
	a, b := pair(t, Config{}, func(a, b *Peer) {
		handle(t, b, "wait", func(ctx context.Context, _ json.RawMessage) (any, error) {
			close(entered)
			<-ctx.Done()
			return nil, ctx.Err()
		})
	})
	raw := ws.Message{Type: ws.Text, Data: []byte(`{"v":1,"type":"call","id":"same","method":"wait","params":null}`)}
	if err := b.receive(raw); err != nil {
		t.Fatal(err)
	}
	<-entered
	if err := b.receive(raw); !errors.Is(err, ErrProtocol) {
		t.Fatal("duplicate active ID accepted", err)
	}
	a.Stop()
	b.Stop()
}