package quic

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"

	"github.com/quic-go/quic-go/quicvarint"
)

var ErrNotInitialPacket = errors.New("not initial packet")

// The Header represents a QUIC header.
type Header struct {
	Type             uint8
	Version          uint32
	SrcConnectionID  []byte
	DestConnectionID []byte
	Length           int64
	Token            []byte
}

// ParseInitialHeader parses the initial packet of a QUIC connection,
// return the initial header and number of bytes read so far.
func ParseInitialHeader(data []byte) (*Header, int64, error) {
	br := bytes.NewReader(data)
	hdr, err := parseLongHeader(br)
	if err != nil {
		return nil, 0, err
	}
	n := int64(len(data) - br.Len())
	return hdr, n, nil
}

func parseLongHeader(b *bytes.Reader) (*Header, error) {
	typeByte, err := b.ReadByte()
	if err != nil {
		return nil, err
	}
	if !isLongHeader(typeByte) {
		return nil, ErrNotInitialPacket
	}
	h := &Header{}
	ver, err := beUint32(b)
	if err != nil {
		return nil, err
	}
	h.Version = ver
	destConnIDLen, err := b.ReadByte()
	if err != nil {
		return nil, err
	}
	h.DestConnectionID = make([]byte, int(destConnIDLen))
	if err := readConnectionID(b, h.DestConnectionID); err != nil {
		return nil, err
	}
	srcConnIDLen, err := b.ReadByte()
	if err != nil {
		return nil, err
	}
	h.SrcConnectionID = make([]byte, int(srcConnIDLen))
	if err := readConnectionID(b, h.SrcConnectionID); err != nil {
		return nil, err
	}

	initialPacketType := byte(0b00)
	if h.Version == V2 {
		initialPacketType = 0b01
	}
	if (typeByte>>4)&0b11 != initialPacketType {
		return nil, ErrNotInitialPacket
	}
	tokenLen, err := quicvarint.Read(b)
	if err != nil {
		return nil, err
	}
	if tokenLen > uint64(b.Len()) {
		return nil, io.EOF
	}
	h.Token = make([]byte, tokenLen)
	if _, err := io.ReadFull(b, h.Token); err != nil {
		return nil, err
	}

	pl, err := quicvarint.Read(b)
	if err != nil {
		return nil, err
	}
	h.Length = int64(pl)
	return h, err
}

func readConnectionID(r io.Reader, cid []byte) error {
	_, err := io.ReadFull(r, cid)
	if err == io.ErrUnexpectedEOF {
		return io.EOF
	}
	return nil
}

func beUint32(r io.Reader) (uint32, error) {
	b := make([]byte, 4)
	if _, err := io.ReadFull(r, b); err != nil {
		return 0, err
	}
	return binary.BigEndian.Uint32(b), nil
}
