package hopp

import "io"
import "net"
import "context"
import "git.tebibyte.media/sashakoshka/hopp/tape"

// TODO: change size limit to be int64

// B implements METADAPT-B over a multiplexed stream-oriented transport such as
// QUIC.
type b struct {
	sizeLimit  int
	underlying MultiConn
}

// AdaptB returns a connection implementing METADAPT-B over a multiplexed
// stream-oriented transport such as QUIC.
func AdaptB(underlying MultiConn) Conn {
	return &b {
		sizeLimit:  defaultSizeLimit,
		underlying: underlying,
	}
}

func (this *b) Close() error {
	return this.underlying.Close()
}

func (this *b) LocalAddr() net.Addr {
	return this.underlying.LocalAddr()
}

func (this *b) RemoteAddr() net.Addr {
	return this.underlying.RemoteAddr()
}

func (this *b) OpenTrans() (Trans, error) {
	stream, err := this.underlying.OpenStream()
	if err != nil { return nil, err }
	return this.newTrans(stream), nil
}

func (this *b) AcceptTrans() (Trans, error) {
	stream, err := this.underlying.AcceptStream(context.Background())
	if err != nil { return nil, err }
	return this.newTrans(stream), nil
}

func (this *b) SetSizeLimit(limit int) {
	this.sizeLimit = limit
}

func (this *b) newTrans(underlying Stream) *transB {
	return &transB {
		sizeLimit: this.sizeLimit,
		underlying: underlying,
	}
}

type transB struct {
	sizeLimit int
	underlying Stream
	currentData io.Reader
}

func (this *transB) Close() error {
	return this.underlying.Close()
}

func (this *transB) ID() int64 {
	return this.underlying.ID()
}

func (this *transB) Send(method uint16, data []byte) error {
	return encodeMessageB(this.underlying, this.sizeLimit, method, data)
}

func (this *transB) Receive() (uint16, []byte, error) {
	// get a reader for the next message
	method, size, data, err := this.receiveReader()
	if err != nil { return 0, nil, err }
	// read the entire thing
	payloadBuffer := make([]byte, int(size))
	_, err = io.ReadFull(data, payloadBuffer)
	if err != nil { return 0, nil, err }
	// we have used up the reader by now so we can forget it exists
	this.currentData = nil
	return method, payloadBuffer, nil
}

func (this *transB) ReceiveReader() (uint16, io.Reader, error) {
	method, _, data, err := this.receiveReader()
	return method, data, err
}

func (this *transB) receiveReader() (uint16, int64, io.Reader, error) {
	// decode the message
	method, size, data, err := decodeMessageB(this.underlying, this.sizeLimit)
	if err != nil { return 0, 0, nil, err }
	// discard current reader if there is one
	if this.currentData == nil {
		io.Copy(io.Discard, this.currentData)
	}
	this.currentData = data
	return method, size, data, nil
}

// MultiConn represens a multiplexed stream-oriented transport for use in
// [AdaptB].
type MultiConn interface {
	// See documentation for [net.Conn].
	io.Closer
	LocalAddr() net.Addr
	RemoteAddr() net.Addr
	// AcceptStream accepts the next incoming stream from the other party.
	AcceptStream(context.Context) (Stream, error)
	// OpenStream opens a new stream.
	OpenStream() (Stream, error)
}

// Stream represents a single stream returned by a [MultiConn].
type Stream interface {
	// See documentation for [net.Conn].
	io.ReadWriteCloser
	// ID returns the stream ID
	ID() int64
}

func encodeMessageB(writer io.Writer, sizeLimit int, method uint16, data []byte) error {
	if len(data) > sizeLimit {
		return ErrPayloadTooLarge
	}
	buffer := make([]byte, 10 + len(data))
	tape.EncodeI16(buffer[:2], method)
	tape.EncodeI64(buffer[2:10], uint64(len(data)))
	copy(buffer[10:], data)
	_, err := writer.Write(buffer)
	return err
}

func decodeMessageB(
	reader io.Reader,
	sizeLimit int,
) (
	method uint16,
	size   int64,
	data   io.Reader,
	err    error,
) {
	headerBuffer := [10]byte { }
	_, err = io.ReadFull(reader, headerBuffer[:])
	if err != nil { return 0, 0, nil, err }
	method, err = tape.DecodeI16[uint16](headerBuffer[:2])
	if err != nil { return 0, 0, nil, err }
	length, err := tape.DecodeI64[uint64](headerBuffer[2:10])
	if err != nil { return 0, 0, nil, err }
	if length > uint64(sizeLimit) {
		return 0, 0, nil, ErrPayloadTooLarge
	}
	return method, int64(length), &io.LimitedReader {
		R: reader,
		N: int64(length),
	}, nil
}