package grpctrace

import (
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

type stream struct {
	ctx    context.Context
	stream grpc.Stream

	recvErr func(error)
}

func (s *stream) Context() context.Context    { return s.ctx }
func (s *stream) SendMsg(m interface{}) error { return s.stream.SendMsg(m) }
func (s *stream) RecvMsg(m interface{}) error {
	err := s.stream.RecvMsg(m)
	if err != nil && s.recvErr != nil {
		s.recvErr(err)
	}
	return err
}

type serverStream struct {
	stream
	server grpc.ServerStream
}

func (ss *serverStream) SendHeader(md metadata.MD) error { return ss.server.SendHeader(md) }
func (ss *serverStream) SetTrailer(md metadata.MD)       { ss.server.SetTrailer(md) }

type clientStream struct {
	stream
	client grpc.ClientStream
}

func (cs *clientStream) Header() (metadata.MD, error) { return cs.client.Header() }
func (cs *clientStream) Trailer() metadata.MD         { return cs.client.Trailer() }
func (cs *clientStream) CloseSend() error             { return cs.client.CloseSend() }
