package stubclient

import (
	"context"
	"net"
	"testing"

	"google.golang.org/grpc"

	stubv1 "github.com/langgenius/dify/dify-agent-runtime/gen/dify/agent/stub/v1"
)

// fakeServer implements the AgentStubService for testing.
type fakeServer struct {
	stubv1.UnimplementedAgentStubServiceServer
}

func (f *fakeServer) Connect(_ context.Context, req *stubv1.ConnectRequest) (*stubv1.ConnectResponse, error) {
	return &stubv1.ConnectResponse{
		ConnectionId: "conn-123",
		Status:       "connected",
	}, nil
}

func (f *fakeServer) CreateFileUploadRequest(_ context.Context, req *stubv1.FileUploadRequest) (*stubv1.FileUploadResponse, error) {
	return &stubv1.FileUploadResponse{
		UploadUrl: "https://upload.example.com/" + req.GetFilename(),
	}, nil
}

func (f *fakeServer) CreateFileDownloadRequest(_ context.Context, req *stubv1.FileDownloadRequest) (*stubv1.FileDownloadResponse, error) {
	return &stubv1.FileDownloadResponse{
		Filename:    "downloaded.txt",
		MimeType:    strPtr("text/plain"),
		Size:        1024,
		DownloadUrl: "https://download.example.com/file",
	}, nil
}

func strPtr(s string) *string { return &s }

func startFakeServer(t *testing.T) string {
	t.Helper()
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("failed to listen: %v", err)
	}
	srv := grpc.NewServer()
	stubv1.RegisterAgentStubServiceServer(srv, &fakeServer{})
	go func() { _ = srv.Serve(lis) }()
	t.Cleanup(srv.Stop)
	return lis.Addr().String()
}

func TestConnect(t *testing.T) {
	addr := startFakeServer(t)
	c, err := Dial(addr)
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer func() { _ = c.Close() }()

	result, err := c.Connect(context.Background(), 1, []string{"hello"}, `{"key":"val"}`)
	if err != nil {
		t.Fatalf("Connect failed: %v", err)
	}
	if result.ConnectionID != "conn-123" {
		t.Errorf("expected conn-123, got %s", result.ConnectionID)
	}
	if result.Status != "connected" {
		t.Errorf("expected connected, got %s", result.Status)
	}
}

func TestCreateFileUpload(t *testing.T) {
	addr := startFakeServer(t)
	c, err := Dial(addr)
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer func() { _ = c.Close() }()

	result, err := c.CreateFileUpload(context.Background(), "test.txt", "text/plain")
	if err != nil {
		t.Fatalf("CreateFileUpload failed: %v", err)
	}
	if result.UploadURL != "https://upload.example.com/test.txt" {
		t.Errorf("unexpected upload URL: %s", result.UploadURL)
	}
}

func TestCreateFileDownload(t *testing.T) {
	addr := startFakeServer(t)
	c, err := Dial(addr)
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer func() { _ = c.Close() }()

	ref := "dify-file-ref:abc123"
	result, err := c.CreateFileDownload(context.Background(), "local_file", &ref, nil, false)
	if err != nil {
		t.Fatalf("CreateFileDownload failed: %v", err)
	}
	if result.Filename != "downloaded.txt" {
		t.Errorf("unexpected filename: %s", result.Filename)
	}
	if result.Size != 1024 {
		t.Errorf("unexpected size: %d", result.Size)
	}
	if result.DownloadURL != "https://download.example.com/file" {
		t.Errorf("unexpected download URL: %s", result.DownloadURL)
	}
}
