package dns

import (
	"strings"
)

type packer struct {
	buf   []byte
	names map[string]int
}

func newPacker() *packer {
	return &packer{
		names: make(map[string]int),
	}
}

func (p *packer) writeUint16(v uint16) {
	p.buf = append(p.buf, byte(v>>8), byte(v))
}

func (p *packer) writeUint32(v uint32) {
	p.buf = append(p.buf, byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
}

func (p *packer) writeName(n Name) error {
	labels := SplitDomainName(n.String())
	for i := 0; i < len(labels); i++ {
		suffix := Fqdn(joinLabels(labels[i:]))
		if off, ok := p.names[suffix]; ok {
			for _, label := range labels[:i] {
				p.buf = append(p.buf, byte(len(label)))
				p.buf = append(p.buf, label...)
			}
			p.writeUint16(uint16(off) | 0xC000)
			return nil
		}
	}

	off := len(p.buf)
	p.buf = append(p.buf, n.Data...)

	for i := 0; i < len(labels); i++ {
		suffix := Fqdn(joinLabels(labels[i:]))
		suffixOff := off
		for j := 0; j < i; j++ {
			suffixOff += 1 + len(labels[j])
		}
		if _, exists := p.names[suffix]; !exists {
			p.names[suffix] = suffixOff
		}
	}
	return nil
}

func joinLabels(labels []string) string {
	if len(labels) == 0 {
		return "."
	}
	return strings.Join(labels, ".") + "."
}

func Pack(m *Msg) ([]byte, error) {
	p := newPacker()

	headerOff := len(p.buf)
	p.buf = append(p.buf, make([]byte, 12)...)

	for _, q := range m.Question {
		if err := q.pack(p); err != nil {
			return nil, err
		}
	}
	for _, rr := range m.Answer {
		if err := rr.pack(p); err != nil {
			return nil, err
		}
	}
	for _, rr := range m.Ns {
		if err := rr.pack(p); err != nil {
			return nil, err
		}
	}
	for _, rr := range m.Extra {
		if err := rr.pack(p); err != nil {
			return nil, err
		}
	}

	putUint16(p.buf[headerOff:], m.Header.ID)
	putUint16(p.buf[headerOff+2:], m.Header.Flags)
	putUint16(p.buf[headerOff+4:], uint16(len(m.Question)))
	putUint16(p.buf[headerOff+6:], uint16(len(m.Answer)))
	putUint16(p.buf[headerOff+8:], uint16(len(m.Ns)))
	putUint16(p.buf[headerOff+10:], uint16(len(m.Extra)))

	return p.buf, nil
}

func putUint16(b []byte, v uint16) {
	b[0] = byte(v >> 8)
	b[1] = byte(v)
}