package main

import (
	"bufio"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"sync"
)

// Wire format shared by the control channel (manager <-> control agent) and
// the data channel (source-agent/dest-agent <-> sink/source-stream). Every
// frame is:
//
//	uint32 BE length  (= 1 + len(payload))
//	byte   type
//	[]byte payload    (length-1 bytes)
const maxFrameLen = 256 * 1024 * 1024 // sanity cap, well above block-size + header defaults

type frameType byte

const (
	frameCtrlJSON  frameType = 0x01 // payload: JSON object (control channel)
	frameHashTable frameType = 0x02 // payload: raw concatenated 32-byte hashes
	frameData      frameType = 0x03 // payload: 8-byte index + 32-byte hash + block bytes
	frameAck       frameType = 0x04 // payload: 8-byte index
	frameErr       frameType = 0x05 // payload: 8-byte index + UTF-8 message
	frameDone      frameType = 0x06 // payload: empty
	frameReady     frameType = 0x07 // payload: empty
)

// FrameWriter serializes concurrent writers onto one underlying stream.
type FrameWriter struct {
	mu sync.Mutex
	w  io.Writer
}

func NewFrameWriter(w io.Writer) *FrameWriter { return &FrameWriter{w: w} }

func (fw *FrameWriter) WriteFrame(typ frameType, payload []byte) error {
	fw.mu.Lock()
	defer fw.mu.Unlock()
	var hdr [4]byte
	binary.BigEndian.PutUint32(hdr[:], uint32(1+len(payload)))
	if _, err := fw.w.Write(hdr[:]); err != nil {
		return err
	}
	if _, err := fw.w.Write([]byte{byte(typ)}); err != nil {
		return err
	}
	if len(payload) > 0 {
		if _, err := fw.w.Write(payload); err != nil {
			return err
		}
	}
	return nil
}

func (fw *FrameWriter) WriteJSON(v any) error {
	b, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return fw.WriteFrame(frameCtrlJSON, b)
}

// FrameReader reads frames from a stream. Not safe for concurrent use.
type FrameReader struct {
	r *bufio.Reader
}

func NewFrameReader(r io.Reader) *FrameReader {
	return &FrameReader{r: bufio.NewReaderSize(r, 64*1024)}
}

func (fr *FrameReader) ReadFrame() (frameType, []byte, error) {
	var hdr [4]byte
	if _, err := io.ReadFull(fr.r, hdr[:]); err != nil {
		return 0, nil, err
	}
	n := binary.BigEndian.Uint32(hdr[:])
	if n == 0 || n > maxFrameLen {
		return 0, nil, fmt.Errorf("frame: invalid length %d", n)
	}
	buf := make([]byte, n)
	if _, err := io.ReadFull(fr.r, buf); err != nil {
		return 0, nil, err
	}
	return frameType(buf[0]), buf[1:], nil
}

func encodeDataFrame(index uint64, hash [32]byte, payload []byte) []byte {
	buf := make([]byte, 8+32+len(payload))
	binary.BigEndian.PutUint64(buf[0:8], index)
	copy(buf[8:40], hash[:])
	copy(buf[40:], payload)
	return buf
}

func decodeDataFrame(b []byte) (index uint64, hash [32]byte, payload []byte, err error) {
	if len(b) < 40 {
		return 0, hash, nil, fmt.Errorf("data frame too short: %d bytes", len(b))
	}
	index = binary.BigEndian.Uint64(b[0:8])
	copy(hash[:], b[8:40])
	payload = b[40:]
	return index, hash, payload, nil
}

func encodeIndexFrame(index uint64) []byte {
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, index)
	return buf
}

func decodeIndexFrame(b []byte) (uint64, error) {
	if len(b) < 8 {
		return 0, fmt.Errorf("index frame too short: %d bytes", len(b))
	}
	return binary.BigEndian.Uint64(b[:8]), nil
}

func encodeErrFrame(index uint64, msg string) []byte {
	buf := make([]byte, 8+len(msg))
	binary.BigEndian.PutUint64(buf[0:8], index)
	copy(buf[8:], msg)
	return buf
}

func decodeErrFrame(b []byte) (uint64, string, error) {
	if len(b) < 8 {
		return 0, "", fmt.Errorf("err frame too short: %d bytes", len(b))
	}
	return binary.BigEndian.Uint64(b[:8]), string(b[8:]), nil
}