//go:generate stringer -type=Type -trimprefix Type
//go:generate stringer -type=SubType -trimprefix SubType
//go:generate stringer -type=AEType -trimprefix AE
package tlv

import (
	"bytes"
	"errors"
	"fmt"

	"inet.af/netaddr"
)

type Type uint8

const (
	TypePad1         = Type(0)
	TypePadN         = Type(1)
	TypeAckReq       = Type(2)
	TypeAck          = Type(3)
	TypeHello        = Type(4)
	TypeIHU          = Type(5)
	TypeRouterID     = Type(6)
	TypeNextHop      = Type(7)
	TypeUpdate       = Type(8)
	TypeRouteRequest = Type(9)
	TypeSeqnoRequest = Type(10)
	TypeTSPC         = Type(11)
	TypeHMAC         = Type(12)
)

type TLV interface {
	T() Type
	L() uint8
}

type SubType uint8

func (s SubType) IsMandatory() bool {
	return s >= 128
}

const (
	SubTypePad1         = SubType(0)
	SubTypePadN         = SubType(1)
	SubTypeDiversity    = SubType(2)
	SubTypeTimestamp    = SubType(3)
	SubTypeSourcePrefix = SubType(128)
)

type SubTLV interface {
	T() SubType
	L() uint8
}

type AEType uint8

const (
	AEWildcard  = AEType(0)
	AEIPv4      = AEType(1)
	AEIPv6      = AEType(2)
	AEIPv6LL    = AEType(3)
	AEIPv4oIPv6 = AEType(4)
)

func AEFromPrefix(p netaddr.IPPrefix) AEType {
	return AEFromIP(p.IP())
}
func AEFromIP(p netaddr.IP) AEType {
	switch {
	case p.IsZero():
		return AEWildcard
	case p.Is4():
		return AEIPv4
	case p.IsLinkLocalUnicast():
		return AEIPv6LL
	case p.Is6():
		return AEIPv6
	}
	panic("unknown AE")
}

type Scanner struct {
	buf []byte
	err error
	t   Type
	l   uint8
}

func NewTLVScanner(buf []byte) *Scanner {
	return &Scanner{
		buf: buf,
	}
}

func (ts *Scanner) Reset(buf []byte) {
	*ts = Scanner{}
	ts.buf = buf
}

func (ts *Scanner) Scan() bool {
	if ts.err != nil {
		return false
	}

	if len(ts.buf) <= int(ts.l) {
		return false
	}

	// move beginning of the buffer to the next TLV
	ts.buf = ts.buf[ts.l:]
	ts.t, ts.l = Type(ts.buf[0]), 0

	switch ts.t {
	case TypePad1:
		ts.buf = ts.buf[1:]
	case TypePadN,
		TypeAckReq,
		TypeAck,
		TypeHello,
		TypeIHU,
		TypeRouterID,
		TypeNextHop,
		TypeUpdate,
		TypeRouteRequest,
		TypeSeqnoRequest,
		TypeTSPC,
		TypeHMAC:
		ts.l = ts.buf[1]
		if 2+int(ts.l) > len(ts.buf) {
			ts.err = fmt.Errorf("Invalid length: type %d, length %d, size %d", ts.t, ts.l, len(ts.buf))
			return false
		}
		// move beginning of the buffer behind the header
		ts.buf = ts.buf[2:]
	default:
		panic("unknown type")
	}

	return true
}

func (ts *Scanner) Err() error {
	return ts.err
}

func (ts *Scanner) Raw() Raw {
	return Raw{ts.t, ts.l, ts.buf[:ts.l:ts.l]}
}

type PacketDecoder struct {
	tlvscanner Scanner

	// parser state
	routerID  RouterID
	nexthopv6 netaddr.IP
	nexthopv4 netaddr.IP
	v6        [16]byte
	v4        [4]byte

	tlv    TLV
	subtlv []byte
	err    error
}

func (s *PacketDecoder) Reset(b []byte, nexthop netaddr.IP) {
	*s = PacketDecoder{}
	s.tlvscanner.Reset(b)
	switch {
	case nexthop.Is6():
		s.nexthopv6 = nexthop
	case nexthop.Is4():
		s.nexthopv4 = nexthop
	}
}
func (s *PacketDecoder) Scan() bool {
	if s.err != nil {
		return false
	}

	if !s.tlvscanner.Scan() {
		if err := s.tlvscanner.Err(); err != nil {
			s.err = err
		}
		return false
	}

	raw := s.tlvscanner.Raw()

	switch raw.T() {
	case TypePad1:
		s.tlv, s.subtlv, s.err = Pad1FromBytes(raw.V())
	case TypePadN:
		s.tlv, s.subtlv, s.err = PadNFromBytes(raw.V())
	case TypeAckReq:
		s.tlv, s.subtlv, s.err = AckReqFromBytes(raw.V())
	case TypeAck:
		s.tlv, s.subtlv, s.err = AckFromBytes(raw.V())
	case TypeHello:
		s.tlv, s.subtlv, s.err = HelloFromBytes(raw.V())
	case TypeIHU:
		s.tlv, s.subtlv, s.err = IHUFromBytes(raw.V())
	case TypeRouterID:
		var rid RouterID
		rid, s.subtlv, s.err = RouterIDFromBytes(raw.V())
		if s.err != nil {
			break
		}

		s.tlv = rid
		s.routerID = rid
	case TypeNextHop:
		var nh NextHop
		nh, s.subtlv, s.err = NextHopFromBytes(raw.V())
		if s.err != nil {
			break
		}

		s.tlv = nh

		switch AEFromIP(nh.Address) {
		case 1:
			s.nexthopv4 = nh.Address
		case 2, 3:
			s.nexthopv6 = nh.Address
		}
	case TypeUpdate:
		s.tlv, s.subtlv, s.err = s.UpdateFromBytes(raw.V())
	case TypeRouteRequest:
		s.tlv, s.subtlv, s.err = RouteRequestFromBytes(raw.V())
	case TypeSeqnoRequest:
		s.tlv, s.subtlv, s.err = SeqnoRequestFromBytes(raw.V())
	case TypeTSPC:
		s.tlv = raw
		s.subtlv = nil
	case TypeHMAC:
		s.tlv = raw
		s.subtlv = nil
	default:
		s.err = fmt.Errorf("Unknown TLV Type: %+v", raw)
	}

	return s.err == nil
}
func (s *PacketDecoder) TLV() TLV {
	if s.err != nil {
		return nil
	}
	return s.tlv
}
func (s *PacketDecoder) SubTLV() []byte {
	if s.err != nil {
		return nil
	}
	return s.subtlv
}
func (s *PacketDecoder) Err() error {
	return s.err
}

type Raw struct {
	t Type
	l uint8
	v []byte
}

func (t Raw) T() Type {
	return t.t
}

func (t Raw) L() uint8 {
	return t.l
}

func (t Raw) V() []byte {
	return t.v
}

type ErrTypeMismatch struct {
	want Type
	got  Type
}

func (e ErrTypeMismatch) Error() string {
	return fmt.Sprintf("Error parsing TLV: Wrong type %s, wanted %s", e.got, e.want)
}

type ErrTLVLength struct {
	l    int
	want int
	t    Type
}

func (e ErrTLVLength) Error() string {
	return fmt.Sprintf("Invalid TLV length: %d, Type %q needs at least %d", e.l, e.t, e.want)
}
func assertLengthGreater(b []byte, t Type, l int) error {
	if len(b) < l {
		return ErrTLVLength{len(b), l, t}
	}
	return nil
}

type ErrSubTLVLength struct {
	l    int
	want int
	t    SubType
}

func (e ErrSubTLVLength) Error() string {
	return fmt.Sprintf("Invalid SubTLV length: %d, Type %q needs at least %d", e.l, e.t, e.want)
}
func assertSubLengthGreater(b []byte, t SubType, l int) error {
	if len(b) < l {
		return ErrSubTLVLength{len(b), l, t}
	}
	return nil
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.1
type Pad1 struct{}

func (Pad1) T() Type {
	return TypePad1
}
func (Pad1) L() uint8 {
	return 8
}
func Pad1FromBytes(b []byte) (Pad1, []byte, error) {
	if err := assertLengthGreater(b, TypePad1, 0); err != nil {
		return Pad1{}, b, err
	}
	return Pad1{}, b[1:], nil
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.2
type PadN uint8

func (PadN) T() Type {
	return TypePadN
}
func (p PadN) L() uint8 {
	return 2 + uint8(p) // TODO: overflow?
}
func PadNFromBytes(b []byte) (PadN, []byte, error) {
	if err := assertLengthGreater(b, TypePadN, 0); err != nil {
		return PadN(0), b, err
	}
	return PadN(b[1]), b[1:], nil
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.3
type AckReq struct {
	Opaque   [2]byte
	Interval uint16
}

func (AckReq) T() Type {
	return TypeAckReq
}
func (a AckReq) L() uint8 {
	return 8
}
func AckReqFromBytes(b []byte) (AckReq, []byte, error) {
	if err := assertLengthGreater(b, TypeAckReq, 6); err != nil {
		return AckReq{}, b, err
	}
	var ar AckReq

	ar.Opaque[0] = b[2]
	ar.Opaque[1] = b[3]
	ar.Interval = uint16(b[4])<<8 | uint16(b[5])

	return ar, b[6:], nil
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.4
type Ack struct {
	Opaque uint16
}

func (Ack) T() Type {
	return TypeAck
}
func (Ack) L() uint8 {
	return 4
}
func AckFromBytes(b []byte) (Ack, []byte, error) {
	if err := assertLengthGreater(b, TypeAck, 2); err != nil {
		return Ack{}, b, err
	}

	return Ack{uint16(b[0])<<8 | uint16(b[1])}, b[4:], nil
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.5
type Hello struct {
	Flags    uint16
	Seqno    uint16
	Interval uint16
}

func (Hello) T() Type {
	return TypeHello
}
func (Hello) L() uint8 {
	return 8
}
func HelloFromBytes(b []byte) (Hello, []byte, error) {
	if err := assertLengthGreater(b, TypeHello, 6); err != nil {
		return Hello{}, b, err
	}
	var h Hello

	h.Flags = uint16(b[0])<<8 | uint16(b[1])
	h.Seqno = uint16(b[2])<<8 | uint16(b[3])
	h.Interval = uint16(b[4])<<8 | uint16(b[5])

	return h, b[6:], nil
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.6
type IHU struct {
	// AE uint8
	// reserved uint8
	Rxcost   uint16
	Interval uint16
	Address  netaddr.IP
}

func (IHU) T() Type {
	return TypeIHU
}
func (i IHU) AE() AEType {
	return AEFromIP(i.Address)
}
func (i IHU) L() uint8 {
	switch i.AE() {
	case 0:
		return 6 // Header only
	case 1:
		return 6 + 4 // IPv4 - no compression
	case 2:
		return 6 + 16 // IPv6 - no compression
	case 3:
		return 6 + 8 // IPv6 Link Local
	}
	panic("invalid AE field")
}
func IHUFromBytes(b []byte) (IHU, []byte, error) {
	if err := assertLengthGreater(b, TypeIHU, 6); err != nil {
		return IHU{}, b, err
	}

	var ihu IHU
	var err error

	ae := AEType(b[0])
	ihu.Rxcost = uint16(b[2])<<8 | uint16(b[3])
	ihu.Interval = uint16(b[4])<<8 | uint16(b[5])
	ihu.Address, b, err = ipFromBytes(ae, b[6:])

	return ihu, b, err
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.7
type RouterID [8]byte

var (
	ErrRouterIDZeros = errors.New("Invalid RouterID: consists of all zeros")
	ErrRouterIDOnes  = errors.New("Invalid RouterID: consists of all ones")
)

func RouterIDFromBytes(b []byte) (RouterID, []byte, error) {
	if err := assertLengthGreater(b, TypeRouterID, 10); err != nil {
		return RouterID{}, b, err
	}

	// skip 2 reserved bytes
	var rid RouterID
	copy(rid[:], b[2:])
	b = b[10:]

	if rid == (RouterID{}) {
		return RouterID{}, b, ErrRouterIDZeros
	}
	if rid == (RouterID{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}) {
		return RouterID{}, b, ErrRouterIDOnes
	}

	return rid, b, nil
}
func (RouterID) T() Type {
	return TypeRouterID
}
func (RouterID) L() uint8 {
	return 8
}
func (r RouterID) String() string {
	var buf bytes.Buffer
	buf.Grow(8*3 - 1)
	for i, b := range r {
		if i > 0 {
			buf.WriteByte(':')
		}
		fmt.Fprintf(&buf, "%02x", b)
	}
	return buf.String()
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.8
type NextHop struct {
	// AE uint8
	// reserved uint8
	Address netaddr.IP
}

func (NextHop) T() Type {
	return TypeNextHop
}
func (n NextHop) L() uint8 {
	switch AEFromIP(n.Address) {
	case 1:
		return 2 + 4 // IPv4 - no compression
	case 2:
		return 2 + 16 // IPv6 - no compression
	case 3:
		return 2 + 8 // IPv6 Link Local
	}
	panic("invalid AE field")
}
func NextHopFromBytes(b []byte) (NextHop, []byte, error) {
	if err := assertLengthGreater(b, TypeNextHop, 6); err != nil {
		return NextHop{}, b, err
	}

	var nh NextHop
	var err error
	ae := AEType(b[0])
	nh.Address, b, err = ipFromBytes(ae, b[2:])
	return nh, b, err
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.9
type Update struct {
	// AE       uint8
	Flags uint8
	// Plen     uint8
	Omitted  uint8
	Interval uint16
	Seqno    uint16
	Metric   uint16
	RouterID RouterID
	Prefix   netaddr.IPPrefix
	NextHop  netaddr.IP
}

func (Update) T() Type {
	return TypeUpdate
}
func (u Update) L() uint8 {
	return 10 + psizeFromPlen(u.Prefix.Bits())
}
func (u Update) FormatHeader() string {
	return fmt.Sprintf("Flags 0x%02x Omitted %2d Interval %4d Seqno %5d Metric %5d RouterID %s",
		u.Flags, u.Omitted, u.Interval, u.Seqno, u.Metric, u.RouterID,
	)
}
func (s *PacketDecoder) UpdateFromBytes(b []byte) (Update, []byte, error) {
	if err := assertLengthGreater(b, TypeUpdate, 10); err != nil {
		return Update{}, b, err
	}
	var u Update
	var err error

	ae := b[0]
	u.Flags = b[1]
	plen := b[2]
	u.Omitted = b[3]
	u.Interval = uint16(b[4])<<8 | uint16(b[5])
	u.Seqno = uint16(b[6])<<8 | uint16(b[7])
	u.Metric = uint16(b[8])<<8 | uint16(b[8])

	b = b[10:]
	switch ae {
	case 0:
	case 1:
		u.Prefix, b, err = prefixV4Default(s.v4, plen, u.Omitted, b)
		if u.Flags&0x80 > 0 {
			s.v4 = u.Prefix.IP().As4()
		}
		u.NextHop = s.nexthopv4
	case 2:
		u.Prefix, b, err = prefixV6Default(s.v6, plen, u.Omitted, b)
		if u.Flags&0x80 > 0 {
			s.v6 = u.Prefix.IP().As16()
		}
		u.NextHop = s.nexthopv6
	case 3:
		u.Prefix, b, err = prefixV6LL(b)
		u.NextHop = s.nexthopv6
	}

	if u.Flags&0x40 > 0 {
		addr := u.Prefix.IP()
		switch {
		case addr.Is6():
			v6 := addr.As16()
			copy(s.routerID[:], v6[8:])
		case addr.Is4():
			s.routerID = RouterID{}
			v4 := addr.As4()
			copy(s.routerID[4:], v4[:])
		}
	}

	u.RouterID = s.routerID

	return u, b, err
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.10
type RouteRequest struct {
	// AE     uint8
	// plen   uint8
	Prefix netaddr.IPPrefix
}

func (r RouteRequest) AE() AEType {
	return AEFromPrefix(r.Prefix)
}
func (RouteRequest) T() Type {
	return TypeRouteRequest
}
func (r RouteRequest) L() uint8 {
	return 4 + psizeFromPlen(r.Prefix.Bits())
}
func RouteRequestFromBytes(b []byte) (RouteRequest, []byte, error) {
	if err := assertLengthGreater(b, TypeRouteRequest, 2); err != nil {
		return RouteRequest{}, b, err
	}
	var rr RouteRequest
	var err error

	rr.Prefix, b, err = prefixUncompressed(b[0], b[1], b[2:])

	return rr, b, err
}

// https://datatracker.ietf.org/doc/html/rfc8966#section-4.6.11
type SeqnoRequest struct {
	// AE    uint8
	// Plen  uint8
	Seqno    uint16
	HopCount uint8
	RouterID RouterID
	Prefix   netaddr.IPPrefix
}

func (SeqnoRequest) T() Type {
	return TypeSeqnoRequest
}
func (r SeqnoRequest) L() uint8 {
	return 14 + psizeFromPlen(r.Prefix.Bits())
}
func SeqnoRequestFromBytes(b []byte) (SeqnoRequest, []byte, error) {
	if err := assertLengthGreater(b, TypeSeqnoRequest, 14); err != nil {
		return SeqnoRequest{}, b, err
	}

	var err error
	var sr SeqnoRequest

	ae := b[0]
	plen := b[1]
	sr.Seqno = uint16(b[2])<<8 | uint16(b[3])
	sr.HopCount = b[4]
	copy(sr.RouterID[:], b[6:])
	sr.Prefix, b, err = prefixUncompressed(ae, plen, b[14:])

	return sr, b, err
}

// https://datatracker.ietf.org/doc/html/draft-ietf-babel-source-specific-07#section-7.1
type SourcePrefix struct {
	netaddr.IPPrefix
}

func (SourcePrefix) T() SubType {
	return SubTypeSourcePrefix
}
func (s SourcePrefix) L() uint8 {
	return psizeFromPlen(s.Bits())
}
func SourcePrefixFromBytes(b []byte) (SourcePrefix, []byte, error) {
	if err := assertSubLengthGreater(b, SubTypeSourcePrefix, 2); err != nil {
		return SourcePrefix{}, b, err
	}

	plen := b[0] // TODO: error on 0
	b = b[1:]
	pfx, b, err := prefixUncompressed(2, plen, b)

	return SourcePrefix{pfx}, b, err
}

func psizeFromPlen(plen uint8) uint8 {
	return uint8((uint16(plen) + 7) / 8)
}
func prefixUncompressed(ae uint8, plen uint8, b []byte) (netaddr.IPPrefix, []byte, error) {
	if len(b) < int(psizeFromPlen(plen)) {
		return netaddr.IPPrefix{}, b, fmt.Errorf("plen too large")
	}

	switch ae {
	case 0:
		return prefixWildcard()
	case 1:
		return prefixV4(plen, b)
	case 2:
		return prefixV6(plen, b)
	case 3:
		return prefixV6LL(b)
	default:
		return netaddr.IPPrefix{}, b, fmt.Errorf("Invalid AE %d", ae)
	}
}

func prefixWildcard() (netaddr.IPPrefix, []byte, error) {
	return netaddr.IPPrefix{}, nil, nil
}

func prefixV4(plen uint8, b []byte) (netaddr.IPPrefix, []byte, error) {
	return prefixV4Default([4]byte{}, plen, 0, b)
}
func prefixV4Default(ip4default [4]byte, plen uint8, omit uint8, b []byte) (netaddr.IPPrefix, []byte, error) {
	var ip4 [4]byte
	psize := psizeFromPlen(plen) - omit
	copy(ip4[:], ip4default[:omit])
	copy(ip4[omit:], b[:psize])
	return netaddr.IPPrefixFrom(netaddr.IPFrom4(ip4), plen), b[psize:], nil
}
func prefixV6(plen uint8, b []byte) (netaddr.IPPrefix, []byte, error) {
	return prefixV6Default([16]byte{}, plen, 0, b)
}
func prefixV6Default(ip6default [16]byte, plen uint8, omit uint8, b []byte) (netaddr.IPPrefix, []byte, error) {
	var ip6 [16]byte
	psize := psizeFromPlen(plen) - omit
	copy(ip6[:], ip6default[:omit])
	copy(ip6[omit:], b[:psize])
	return netaddr.IPPrefixFrom(netaddr.IPv6Raw(ip6), plen), b[psize:], nil
}
func prefixV6LL(b []byte) (netaddr.IPPrefix, []byte, error) {
	var ip6ll [16]byte
	ip6ll[0] = 0xfe
	ip6ll[1] = 0x80
	copy(ip6ll[8:], b[:8])
	return netaddr.IPPrefixFrom(netaddr.IPv6Raw(ip6ll), 8), b[8:], nil
}
func ipFromBytes(ae AEType, b []byte) (netaddr.IP, []byte, error) {
	switch ae {
	case AEWildcard:
		return netaddr.IP{}, b, nil
	case AEIPv4:
		if len(b) < 4 {
			return netaddr.IP{}, b, fmt.Errorf("Not enough bytes for v4 address: %d", len(b))
		}
		var ip4 [4]byte
		copy(ip4[:], b[:4])
		return netaddr.IPFrom4(ip4), b[4:], nil
	case AEIPv6:
		if len(b) < 16 {
			return netaddr.IP{}, b, fmt.Errorf("Not enough bytes for v6 address: %d", len(b))
		}
		var ip6 [16]byte
		copy(ip6[:], b[:16])
		return netaddr.IPv6Raw(ip6), b[16:], nil
	case AEIPv6LL:
		if len(b) < 8 {
			return netaddr.IP{}, b, fmt.Errorf("Not enough bytes for v6ll address: %d", len(b))
		}
		var ip6ll [16]byte
		ip6ll[0] = 0xfe
		ip6ll[1] = 0x80
		copy(ip6ll[8:], b[:8])
		return netaddr.IPv6Raw(ip6ll), b[8:], nil
	case AEIPv4oIPv6:
		return netaddr.IP{}, b, fmt.Errorf("Not implemented AE Type %s", ae)
	default:
		return netaddr.IP{}, b, fmt.Errorf("Invalid AE %d", ae)
	}
}