package bt

// https://wiki.theory.org/BitTorrentSpecification#Messages
// Messages: <length prefix><message ID><payload>
//   keep-alive: <len=0000>
//   choke: <len=0001><type=0>
//   unchoke: <len=0001><type=1>
//   interested: <len=0001><type=2>
//   not interested: <len=0001><type=3>
//   have: <len=0005><type=4><piece index>
//   bitfield: <len=0001+X><type=5><bitfield>
//   request: <len=0013><type=6><index><begin><length>
//   piece: <len=0009+X><type=7><index><begin><block>
//   cancel: <len=0013><type=8><index><begin><length>

import (
	"encoding/binary"
	"io"
)

type MsgKind uint8

const BlockSize = 16384

const (
	Choke MsgKind = iota
	Unchoke
	Interested
	NotInterested
	Have
	Bitfield
	Request
	Piece
	Cancel
)

type Msg struct {
	Kind     MsgKind
	Index    uint32
	Begin    uint32
	Length   uint32
	Bitfield []byte
	Block    []byte
}

// [4: len] [1: type] [len-1: payload]
func readMsg(r io.Reader) (*Msg, int, error) {
	var lenBuf [4]byte
	if _, err := io.ReadFull(r, lenBuf[:]); err != nil {
		return nil, 0, err
	}
	n := binary.BigEndian.Uint32(lenBuf[:])
	if n == 0 {
		return nil, 4, nil
	}
	buf := make([]byte, n)
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, 0, err
	}
	msg := &Msg{Kind: MsgKind(buf[0])}
	payload := buf[1:]
	switch msg.Kind {
	case Have:
		msg.Index = binary.BigEndian.Uint32(payload)
	case Bitfield:
		msg.Bitfield = payload
	case Request, Cancel:
		msg.Index = binary.BigEndian.Uint32(payload[0:4])
		msg.Begin = binary.BigEndian.Uint32(payload[4:8])
		msg.Length = binary.BigEndian.Uint32(payload[8:12])
	case Piece:
		msg.Index = binary.BigEndian.Uint32(payload[0:4])
		msg.Begin = binary.BigEndian.Uint32(payload[4:8])
		msg.Block = payload[8:]
	}
	return msg, int(4 + n), nil
}

func writeMsg(w io.Writer, msg *Msg) (int, error) {
	if msg == nil {
		return w.Write(make([]byte, 4))
	}

	var payload []byte
	switch msg.Kind {
	case Choke, Unchoke, Interested, NotInterested:
	case Have:
		payload = make([]byte, 4)
		binary.BigEndian.PutUint32(payload, msg.Index)
	case Bitfield:
		payload = msg.Bitfield
	case Request, Cancel:
		payload = make([]byte, 12)
		binary.BigEndian.PutUint32(payload[0:4], msg.Index)
		binary.BigEndian.PutUint32(payload[4:8], msg.Begin)
		binary.BigEndian.PutUint32(payload[8:12], msg.Length)
	case Piece:
		payload = make([]byte, 8+len(msg.Block))
		binary.BigEndian.PutUint32(payload[0:4], msg.Index)
		binary.BigEndian.PutUint32(payload[4:8], msg.Begin)
		copy(payload[8:], msg.Block)
	}
	buf := make([]byte, 5+len(payload))
	binary.BigEndian.PutUint32(buf[0:4], uint32(1+len(payload)))
	buf[4] = byte(msg.Kind)
	copy(buf[5:], payload)
	return w.Write(buf)
}

func SetBit(bf []byte, i int) {
	bf[i/8] |= 1 << (7 - i%8)
}

func HasBit(bf []byte, i int) bool {
	return bf[i/8]&(1<<(7-i%8)) != 0
}