// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package webtransport

import (
	"bytes"
	"sync"
	"time"
)

// PreparedMessage caches on the wire representations of a message payload.
// Use PreparedMessage to efficiently send a message payload to multiple
// connections. PreparedMessage is especially useful when compression is used
// because the CPU and memory expensive compression operation can be executed
// once for a given set of compression options.
type PreparedMessage struct {
	messageType int
	data        []byte
	mu          sync.Mutex
	frames      map[prepareKey]*preparedFrame
}

// prepareKey defines a unique set of options to cache prepared frames in PreparedMessage.
type prepareKey struct {
	isServer bool
}

// preparedFrame contains data in wire representation.
type preparedFrame struct {
	once sync.Once
	data []byte
}

// NewPreparedMessage returns an initialized PreparedMessage. You can then send
// it to connection using WritePreparedMessage method. Valid wire
// representation will be calculated lazily only once for a set of current
// connection options.
func NewPreparedMessage(messageType int, data []byte) (*PreparedMessage, error) {
	pm := &PreparedMessage{
		messageType: messageType,
		frames:      make(map[prepareKey]*preparedFrame),
		data:        data,
	}

	// Prepare a plain server frame.
	_, frameData, err := pm.frame(prepareKey{isServer: true})
	if err != nil {
		return nil, err
	}

	// To protect against caller modifying the data argument, remember the data
	// copied to the plain server frame.
	pm.data = frameData[len(frameData)-len(data):]
	return pm, nil
}

func (pm *PreparedMessage) frame(key prepareKey) (int, []byte, error) {
	pm.mu.Lock()
	frame, ok := pm.frames[key]
	if !ok {
		frame = &preparedFrame{}
		pm.frames[key] = frame
	}
	pm.mu.Unlock()

	var err error
	frame.once.Do(func() {
		// Prepare a frame using a 'fake' connection.
		// TODO: Refactor code in conn.go to allow more direct construction of
		// the frame.
		mu := make(chan struct{}, 1)
		mu <- struct{}{}
		var nc prepareConn
		c := &Conn{
			stream:   &nc,
			mu:       mu,
			isServer: key.isServer,
			writeBuf: make([]byte, defaultWriteBufferSize+maxFrameHeaderSize),
		}
		err = c.WriteMessage(pm.messageType, pm.data)
		frame.data = nc.buf.Bytes()
	})
	return pm.messageType, frame.data, err
}

type prepareConn struct {
	buf bytes.Buffer
}

var (
	_ streamWithDeadline = &prepareConn{}
)

func (pc *prepareConn) Read(p []byte) (n int, err error)   { return pc.buf.Read(p) }
func (pc *prepareConn) Write(p []byte) (int, error)        { return pc.buf.Write(p) }
func (pc *prepareConn) SetWriteDeadline(t time.Time) error { return nil }
func (pc *prepareConn) SetReadDeadline(t time.Time) error  { return nil }
func (pc *prepareConn) Close() error                       { return nil }