package iskra

import "io"

// Op codes for the iskra wire protocol.
const (
	OpBigramWeight        uint8 = 0x01
	OpBigramWeightRelaxed uint8 = 0x02
	OpIngestBigram        uint8 = 0x03
	OpIngestText          uint8 = 0x04
	OpWalkStep            uint8 = 0x05
	OpTranslate           uint8 = 0x06
	OpLookup              uint8 = 0x07
	OpInsert              uint8 = 0x08
	OpBeamRun             uint8 = 0x09
	OpFlush               uint8 = 0x0A
	OpIngestRef           uint8 = 0x0B
	OpGetRefs             uint8 = 0x0C
	OpWalkOutward         uint8 = 0x0D
)

// Response status codes.
const (
	StatusOK       uint8 = 0
	StatusNotFound uint8 = 1
	StatusError    uint8 = 2
)

// Frame layout: [length:4 LE][op:1][payload:N]
// Response layout: [length:4 LE][status:1][payload:N]

// WriteFrame writes a framed message (length-prefixed) to w.
func WriteFrame(w io.Writer, data []byte) error {
	var hdr [4]byte
	n := uint32(len(data))
	hdr[0] = byte(n)
	hdr[1] = byte(n >> 8)
	hdr[2] = byte(n >> 16)
	hdr[3] = byte(n >> 24)
	if _, err := w.Write(hdr[:]); err != nil {
		return err
	}
	_, err := w.Write(data)
	return err
}

// ReadFrame reads a length-prefixed frame from r.
// Returns nil on EOF or error.
func ReadFrame(r io.Reader) ([]byte, error) {
	var hdr [4]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return nil, err
	}
	n := uint32(hdr[0]) | uint32(hdr[1])<<8 | uint32(hdr[2])<<16 | uint32(hdr[3])<<24
	if n == 0 {
		return []byte{}, nil
	}
	if n > 4<<20 {
		return nil, io.ErrUnexpectedEOF
	}
	buf := []byte{:int32(n)}
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, err
	}
	return buf, nil
}

// EncodeString appends a length-prefixed string (uint16 LE + bytes).
func EncodeString(buf []byte, s string) []byte {
	n := uint16(len(s))
	buf = append(buf, byte(n), byte(n>>8))
	buf = append(buf, []byte(s)...)
	return buf
}

// DecodeString reads a length-prefixed string from buf at offset.
// Returns (string, bytesConsumed). Returns ("", 0) on underflow.
func DecodeString(buf []byte, off int32) (string, int32) {
	if off+2 > len(buf) {
		return "", 0
	}
	n := int32(uint16(buf[off]) | uint16(buf[off+1])<<8)
	off += 2
	if off+n > len(buf) {
		return "", 0
	}
	return string(buf[off : off+n]), 2 + n
}

// EncodeUint32 appends a uint32 LE.
func EncodeUint32(buf []byte, v uint32) []byte {
	return append(buf, byte(v), byte(v>>8), byte(v>>16), byte(v>>24))
}

// DecodeUint32 reads a uint32 LE from buf at offset.
func DecodeUint32(buf []byte, off int32) (uint32, int32) {
	if off+4 > len(buf) {
		return 0, 0
	}
	v := uint32(buf[off]) | uint32(buf[off+1])<<8 | uint32(buf[off+2])<<16 | uint32(buf[off+3])<<24
	return v, 4
}

// EncodeUint64 appends a uint64 LE.
func EncodeUint64(buf []byte, v uint64) []byte {
	return append(buf, byte(v), byte(v>>8), byte(v>>16), byte(v>>24),
		byte(v>>32), byte(v>>40), byte(v>>48), byte(v>>56))
}

// DecodeUint64 reads a uint64 LE from buf at offset.
func DecodeUint64(buf []byte, off int32) (uint64, int32) {
	if off+8 > len(buf) {
		return 0, 0
	}
	v := uint64(buf[off]) | uint64(buf[off+1])<<8 | uint64(buf[off+2])<<16 | uint64(buf[off+3])<<24 |
		uint64(buf[off+4])<<32 | uint64(buf[off+5])<<40 | uint64(buf[off+6])<<48 | uint64(buf[off+7])<<56
	return v, 8
}

// Request encoders

func EncodeBigramWeightReq(domain uint8, coord uint64, prev, next string) []byte {
	buf := []byte{:0:3 + 8 + len(prev) + len(next)}
	buf = append(buf, OpBigramWeight, domain)
	buf = EncodeUint64(buf, coord)
	buf = EncodeString(buf, prev)
	buf = EncodeString(buf, next)
	return buf
}

func EncodeBigramWeightRelaxedReq(domain uint8, coord uint64, prev, next string) []byte {
	buf := []byte{:0:3 + 8 + len(prev) + len(next)}
	buf = append(buf, OpBigramWeightRelaxed, domain)
	buf = EncodeUint64(buf, coord)
	buf = EncodeString(buf, prev)
	buf = EncodeString(buf, next)
	return buf
}

func EncodeIngestBigramReq(domain uint8, coord uint64, prev, next string) []byte {
	buf := []byte{:0:3 + 8 + len(prev) + len(next)}
	buf = append(buf, OpIngestBigram, domain)
	buf = EncodeUint64(buf, coord)
	buf = EncodeString(buf, prev)
	buf = EncodeString(buf, next)
	return buf
}

func EncodeIngestTextReq(domain uint8, coord uint64, tokens []string) []byte {
	total := 2 + 8 + 2
	for _, tok := range tokens {
		total += 2 + len(tok)
	}
	buf := []byte{:0:total}
	buf = append(buf, OpIngestText, domain)
	buf = EncodeUint64(buf, coord)
	buf = append(buf, byte(len(tokens)), byte(len(tokens)>>8))
	for _, tok := range tokens {
		buf = EncodeString(buf, tok)
	}
	return buf
}

func EncodeWalkStepReq(domain uint8, coord uint64, word string, maxCandidates uint16) []byte {
	buf := []byte{:0:2 + 8 + 2 + len(word) + 2}
	buf = append(buf, OpWalkStep, domain)
	buf = EncodeUint64(buf, coord)
	buf = EncodeString(buf, word)
	buf = append(buf, byte(maxCandidates), byte(maxCandidates>>8))
	return buf
}

func EncodeTranslateReq(srcDomain, dstDomain uint8, coord uint64, word string) []byte {
	buf := []byte{:0:3 + 8 + 2 + len(word)}
	buf = append(buf, OpTranslate, srcDomain, dstDomain)
	buf = EncodeUint64(buf, coord)
	buf = EncodeString(buf, word)
	return buf
}

func EncodeBeamRunReq(domain uint8, coord uint64, word string, width, maxSteps uint16, kind WalkKind) []byte {
	buf := []byte{:0:2 + 8 + 2 + len(word) + 5}
	buf = append(buf, OpBeamRun, domain)
	buf = EncodeUint64(buf, coord)
	buf = EncodeString(buf, word)
	buf = append(buf, byte(width), byte(width>>8))
	buf = append(buf, byte(maxSteps), byte(maxSteps>>8))
	buf = append(buf, byte(kind))
	return buf
}

func EncodeFlushReq() []byte {
	return []byte{OpFlush}
}

func EncodeIngestRefReq(domain uint8, coord uint64, src, target string, kind uint8) []byte {
	buf := []byte{:0:3 + 8 + len(src) + len(target) + 4}
	buf = append(buf, OpIngestRef, domain)
	buf = EncodeUint64(buf, coord)
	buf = EncodeString(buf, src)
	buf = EncodeString(buf, target)
	buf = append(buf, kind)
	return buf
}

func EncodeGetRefsReq(domain uint8, src string) []byte {
	buf := []byte{:0:2 + 2 + len(src)}
	buf = append(buf, OpGetRefs, domain)
	buf = EncodeString(buf, src)
	return buf
}

func EncodeWalkOutwardReq(domain uint8, src string, maxDepth uint8) []byte {
	buf := []byte{:0:3 + 2 + len(src)}
	buf = append(buf, OpWalkOutward, domain)
	buf = EncodeString(buf, src)
	buf = append(buf, maxDepth)
	return buf
}

func EncodeRefsResponse(refs []RefEntry) []byte {
	buf := []byte{:0:3 + len(refs)*16}
	buf = append(buf, StatusOK)
	buf = append(buf, byte(len(refs)), byte(len(refs)>>8))
	for _, r := range refs {
		buf = EncodeUint32(buf, r.RecIdx)
		buf = EncodeString(buf, r.Target)
		buf = append(buf, r.Kind)
		buf = EncodeUint32(buf, r.Weight)
	}
	return buf
}

// Response encoders

func EncodeResponse(status uint8, payload []byte) []byte {
	buf := []byte{:0:1 + len(payload)}
	buf = append(buf, status)
	buf = append(buf, payload...)
	return buf
}

func EncodeUint32Response(v uint32) []byte {
	buf := []byte{StatusOK, 0, 0, 0, 0}
	buf[1] = byte(v)
	buf[2] = byte(v >> 8)
	buf[3] = byte(v >> 16)
	buf[4] = byte(v >> 24)
	return buf
}

func EncodeWeightRelaxedResponse(w uint32, coord uint64) []byte {
	buf := []byte{:0:13}
	buf = append(buf, StatusOK)
	buf = EncodeUint32(buf, w)
	buf = EncodeUint64(buf, coord)
	return buf
}

func EncodeStringResponse(s string) []byte {
	buf := []byte{:0:3 + len(s)}
	buf = append(buf, StatusOK)
	buf = EncodeString(buf, s)
	return buf
}

func EncodeWalkCandidatesResponse(candidates []WalkCandidate) []byte {
	buf := []byte{:0:3 + len(candidates)*32}
	buf = append(buf, StatusOK)
	buf = append(buf, byte(len(candidates)), byte(len(candidates)>>8))
	for _, c := range candidates {
		buf = append(buf, c.State.Domain)
		buf = EncodeUint64(buf, c.State.Coord)
		buf = EncodeString(buf, c.State.Word)
		buf = EncodeUint32(buf, c.State.Score)
		buf = append(buf, c.State.Depth)
		buf = EncodeUint32(buf, c.Weight)
	}
	return buf
}

func EncodeBeamResponse(s WalkState) []byte {
	buf := []byte{:0:16 + len(s.Word)}
	buf = append(buf, StatusOK)
	buf = append(buf, s.Domain, s.Depth)
	buf = EncodeUint64(buf, s.Coord)
	buf = EncodeUint32(buf, s.Score)
	buf = EncodeString(buf, s.Word)
	return buf
}