package protocol

import (
	"encoding/binary"
	"fmt"
	"io"
)

type FrameReadWriter interface {
	ReadFrame() (*Frame, error)
	WriteFrame(opCode byte, payload []byte) error
}

const maxFramePayloadSize = 1 << 24

type Frame struct {
	OpCode  byte
	Payload []byte
}

func EncodeFrame(opCode byte, payload []byte) []byte {
	payloadLen := uint64(len(payload))
	buf := make([]byte, 1+varintLen(payloadLen)+len(payload))
	buf[0] = opCode
	off := 1
	off += encodeVarint(buf[off:], uint64(payloadLen))
	copy(buf[off:], payload)
	return buf
}

func ReadFrame(r io.Reader) (*Frame, error) {
	opBuf := make([]byte, 1)
	if _, err := io.ReadFull(r, opBuf); err != nil {
		return nil, fmt.Errorf("read opcode: %w", err)
	}
	opCode := opBuf[0]

	payloadLen, err := readVarint(r)
	if err != nil {
		return nil, fmt.Errorf("read payload length: %w", err)
	}
	if payloadLen > maxFramePayloadSize {
		return nil, fmt.Errorf("payload too large: %d", payloadLen)
	}

	payload := make([]byte, payloadLen)
	if payloadLen > 0 {
		if _, err := io.ReadFull(r, payload); err != nil {
			return nil, fmt.Errorf("read payload: %w", err)
		}
	}

	return &Frame{OpCode: opCode, Payload: payload}, nil
}

func WriteFrame(w io.Writer, opCode byte, payload []byte) error {
	data := EncodeFrame(opCode, payload)
	_, err := w.Write(data)
	return err
}

func encodeVarint(buf []byte, v uint64) int {
	i := 0
	for v >= 0x80 {
		buf[i] = byte(v) | 0x80
		v >>= 7
		i++
	}
	buf[i] = byte(v)
	return i + 1
}

func varintLen(v uint64) int {
	l := 0
	for {
		l++
		if v < 0x80 {
			break
		}
		v >>= 7
	}
	return l
}

func readVarint(r io.Reader) (uint64, error) {
	var v uint64
	var shift uint
	for {
		b := make([]byte, 1)
		if _, err := io.ReadFull(r, b); err != nil {
			return 0, err
		}
		v |= uint64(b[0]&0x7F) << shift
		if b[0]&0x80 == 0 {
			break
		}
		shift += 7
		if shift > 63 {
			return 0, fmt.Errorf("varint too long")
		}
	}
	return v, nil
}

func ReadUint64(buf []byte) uint64 {
	return binary.BigEndian.Uint64(buf)
}

func PutUint64(buf []byte, v uint64) {
	binary.BigEndian.PutUint64(buf, v)
}

func ReadUint32(buf []byte) uint32 {
	return binary.BigEndian.Uint32(buf)
}

func PutUint32(buf []byte, v uint32) {
	binary.BigEndian.PutUint32(buf, v)
}