package pipe

import (
	"context"
	"errors"
	fmt "fmt"
	io "io"

	"github.com/dagger/dagger/util/grpcutil"
	grpc "google.golang.org/grpc"
	codes "google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/anypb"
)

// PipeAttachable implements the PipeServer interface generated by protoc.
type PipeAttachable struct {
	rootCtx context.Context

	stdin  io.Reader
	stdout io.Writer

	UnimplementedPipeServer
}

func NewPipeAttachable(rootCtx context.Context, stdin io.Reader, stdout io.Writer) PipeAttachable {
	return PipeAttachable{
		rootCtx: rootCtx,
		stdin:   stdin,
		stdout:  stdout,
	}
}

func (p PipeAttachable) Register(srv *grpc.Server) {
	RegisterPipeServer(srv, p)
}

func (p PipeAttachable) IO(srv Pipe_IOServer) error {
	ctx := p.rootCtx
	pio := &PipeIO{GRPC: srv}
	go io.Copy(pio, newCtxReader(ctx, p.stdin))
	_, err := io.Copy(p.stdout, newCtxReader(ctx, pio))
	return err
}

// PipeIO transforms a Pipe_IOServer or a Pipe_IOClient into an io.ReadWriter
type PipeIO struct {
	GRPC interface {
		Send(*Data) error
		Recv() (*Data, error)
	}
	rem []byte // remainder buffer
}

func (pio *PipeIO) Write(p []byte) (n int, err error) {
	err = pio.GRPC.Send(&Data{Data: p})
	if err != nil {
		return 0, fmt.Errorf("error writing dagger pipe: %w", err)
	}
	return len(p), nil
}

func (pio *PipeIO) Read(p []byte) (n int, err error) {
	// read from the remainder buffer first
	n = copy(p, pio.rem)
	p = p[n:]
	pio.rem = pio.rem[n:]
	if len(p) == 0 || n != 0 {
		return n, nil
	}
	req, err := pio.GRPC.Recv()
	if err != nil {
		// Return io.EOF in certain cases to conform to io.Reader
		//
		// The following conditionals were copied and modified manually from the Terminal session attachable,
		// where `nil` was being returned as opposed to `io.EOF` returned here.
		// The reason for discrepancy is because this is an attempt to simplify the logic using io.Copy, io.Reader
		// interfaces which mandate io.EOF here.
		// In the future, terminal session attachable may also get refactored using this simplified logic.
		if errors.Is(err, context.Canceled) || status.Code(err) == codes.Canceled {
			// canceled
			return 0, io.EOF
		}
		if errors.Is(err, io.EOF) {
			// stopped
			return 0, io.EOF
		}
		if status.Code(err) == codes.Unavailable {
			// client disconnected (i.e. quitting Dagger out)
			return 0, io.EOF
		}
		return 0, fmt.Errorf("error reading dagger pipe: %w", err)
	}
	pio.rem = req.GetData()
	n = copy(p, pio.rem)
	pio.rem = pio.rem[n:]
	return n, nil
}

func (pio *PipeIO) Close() error {
	if c, ok := pio.GRPC.(interface {
		CloseSend() error
	}); ok {
		return c.CloseSend()
	}
	return nil
}

type PipeProxy struct {
	client PipeClient
}

func NewPipeProxy(client PipeClient) PipeProxy {
	return PipeProxy{
		client: client,
	}
}

func (p PipeProxy) Register(srv *grpc.Server) {
	RegisterPipeServer(srv, p)
}

func (p PipeProxy) IO(stream Pipe_IOServer) error {
	ctx, cancel := context.WithCancelCause(stream.Context())
	defer cancel(errors.New("proxy stream closed"))

	clientStream, err := p.client.IO(grpcutil.IncomingToOutgoingContext(ctx))
	if err != nil {
		return fmt.Errorf("create client stream: %w", err)
	}

	return grpcutil.ProxyStream[anypb.Any](ctx, clientStream, stream)
}
