package wire

import (
	"bytes"
	"fmt"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
)

// The Header is the header of a QUIC Packet.
// TODO: add support for the key phase
type Header struct {
	Type             uint8
	IsLongHeader     bool
	KeyPhase         int
	OmitConnectionID bool
	ConnectionID     protocol.ConnectionID
	PacketNumber     protocol.PacketNumber
	PacketNumberLen  protocol.PacketNumberLen
	Version          protocol.VersionNumber
}

// ParseHeader parses a header
func ParseHeader(b *bytes.Reader) (*Header, error) {
	typeByte, err := b.ReadByte()
	if err != nil {
		return nil, err
	}
	if typeByte&0x80 > 0 {
		return parseLongHeader(b, typeByte)
	}
	return parseShortHeader(b, typeByte)
}

func parseLongHeader(b *bytes.Reader, typeByte byte) (*Header, error) {
	connID, err := utils.BigEndian.ReadUint64(b)
	if err != nil {
		return nil, err
	}
	pn, err := utils.BigEndian.ReadUint32(b)
	if err != nil {
		return nil, err
	}
	v, err := utils.BigEndian.ReadUint32(b)
	if err != nil {
		return nil, err
	}
	return &Header{
		Type:            typeByte & 0x7f,
		IsLongHeader:    true,
		ConnectionID:    protocol.ConnectionID(connID),
		PacketNumber:    protocol.PacketNumber(pn),
		PacketNumberLen: protocol.PacketNumberLen4,
		Version:         protocol.VersionNumber(v),
	}, nil
}

func parseShortHeader(b *bytes.Reader, typeByte byte) (*Header, error) {
	hasConnID := typeByte&0x40 > 0
	var connID uint64
	if hasConnID {
		var err error
		connID, err = utils.BigEndian.ReadUint64(b)
		if err != nil {
			return nil, err
		}
	}
	pnLen := 1 << ((typeByte & 0x3) - 1)
	pn, err := utils.BigEndian.ReadUintN(b, uint8(pnLen))
	if err != nil {
		return nil, err
	}
	return &Header{
		KeyPhase:         int(typeByte&0x20) >> 5,
		OmitConnectionID: !hasConnID,
		ConnectionID:     protocol.ConnectionID(connID),
		PacketNumber:     protocol.PacketNumber(pn),
		PacketNumberLen:  protocol.PacketNumberLen(pnLen),
	}, nil
}

func (h *Header) Write(b *bytes.Buffer) error {
	if h.IsLongHeader {
		return h.writeLongHeader(b)
	}
	return h.writeShortHeader(b)
}

func (h *Header) writeLongHeader(b *bytes.Buffer) error {
	b.WriteByte(byte(0x80 ^ h.Type))
	utils.BigEndian.WriteUint64(b, uint64(h.ConnectionID))
	utils.BigEndian.WriteUint32(b, uint32(h.PacketNumber))
	utils.BigEndian.WriteUint32(b, uint32(h.Version))
	return nil
}

func (h *Header) writeShortHeader(b *bytes.Buffer) error {
	typeByte := byte(h.KeyPhase << 5)
	if !h.OmitConnectionID {
		typeByte ^= 0x40
	}
	switch h.PacketNumberLen {
	case protocol.PacketNumberLen1:
		typeByte ^= 0x1
	case protocol.PacketNumberLen2:
		typeByte ^= 0x2
	case protocol.PacketNumberLen4:
		typeByte ^= 0x3
	default:
		return fmt.Errorf("invalid packet number length: %d", h.PacketNumberLen)
	}

	b.WriteByte(typeByte)
	if !h.OmitConnectionID {
		utils.BigEndian.WriteUint64(b, uint64(h.ConnectionID))
	}
	switch h.PacketNumberLen {
	case protocol.PacketNumberLen1:
		b.WriteByte(uint8(h.PacketNumber))
	case protocol.PacketNumberLen2:
		utils.BigEndian.WriteUint16(b, uint16(h.PacketNumber))
	case protocol.PacketNumberLen4:
		utils.BigEndian.WriteUint32(b, uint32(h.PacketNumber))
	}
	return nil
}

// GetLength gets the length of the Header in bytes.
func (h *Header) GetLength() (protocol.ByteCount, error) {
	if h.IsLongHeader {
		return 1 + 8 + 4 + 4, nil
	}

	length := protocol.ByteCount(1) // type byte
	if !h.OmitConnectionID {
		length += 8
	}
	if h.PacketNumberLen != protocol.PacketNumberLen1 && h.PacketNumberLen != protocol.PacketNumberLen2 && h.PacketNumberLen != protocol.PacketNumberLen4 {
		return 0, fmt.Errorf("invalid packet number length: %d", h.PacketNumberLen)
	}
	length += protocol.ByteCount(h.PacketNumberLen)
	return length, nil
}